Ra sdxs-VAE (step 61000)

.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

README.md

@@ -0,0 +1,80 @@
+---
+license: mit
+pipeline_tag: text-to-image
+tags: [text-to-image, flow-matching, pyramidal-unet, latent-diffusion, asym-vae, from-scratch, ra, udit]
+---
+
+# Ra2
+
+**Ra2** — second-generation layout of Ra: the transformer no longer carries a pixel-space refiner, so the
+model is a plain latent diffusion model and the image is decoded by a **separate asymmetric VAE**
+(f8 encoder / f16 decoder → 512×512 from a 32×32 latent).
+
+It is a 256→512 text-to-image model trained from scratch: a pyramidal U-transformer (240M, 40 blocks)
+flow-matching in an 8×/32-channel latent, a frozen Ettin-150M text encoder with a learnable layer-slice
+mixer (10.0M), and a **separate asymmetric VAE** that decodes the latent into pixels.
+
+| | |
+|---|---|
+| output | 512×512 from a 32×32 latent (encoder f8, decoder f16) |
+| sampler | 16 Euler steps, cfg 2.5 |
+| latent | 32 channels from the bundled `vae/` (AsymmetricAutoencoderKL) |
+| text encoder | [`jhu-clsp/ettin-encoder-150m`](https://huggingface.co/jhu-clsp/ettin-encoder-150m), frozen, bundled in `text_encoder/` |
+| body | 40 blocks — L0: 6 conv + 8 attention, L1: 9, L2: 9, L3: 4 (widths 192→512→1024→1536) |
+| training | latent flow matching only (MSE in latent space); the VAE decodes, no pixel-space refiner |
+| licence | MIT |
+
+## Design
+
+The body is a U-stack: width grows as tokens shrink, so expensive attention runs on 16–64 tokens while
+the 1024-token grid keeps 3×3 convolutions, and skip connections add every level back in the decoder.
+One L0-attention block costs 0.86 GMAC (3.9% of the body) at 0.37M parameters, which is why extra blocks
+were added at full resolution rather than in depth.
+
+Positions: 2D RoPE inside attention plus an absolute 2D sin/cos embedding projected into the level width
+by a zero-initialised linear layer, so new blocks can be grafted onto a trained model without changing
+its output (verified bit-identical: `Body.graft`).
+
+Decoding is plain: the latent goes through the VAE decoder once. That keeps the objective a clean latent
+MSE — the price is that fine detail must live in the 32×32 latent.
+
+## Training
+
+| stage | what | steps |
+|---|---|---|
+| 1 | FLUX-Reason-6M, 256², old f8c32 AE + pixel-space refiner (23.9M) | 560 142 |
+| 2 | same data, capacity graft 26 → 36 blocks, absolute-PE branch | 404 000 |
+| 3 | switch to the asymmetric VAE (32×32 latent → 512²), latent-only loss, MJ shard | 5 000 |
+| 4 | butterfly fine-tune, 966 images, zero-init graft 36 → 40 blocks (3+3 L0 attention) + EMA | 25 000 |
+
+The bundled `transformer/` is the stage-4 checkpoint (step 61 000 of the current counter). It is a
+**research stage**, not a general-purpose model: it was trained to fit one narrow dataset.
+
+## Examples
+
+The current checkpoint on its training domain (butterfly specimens, 512² downscaled to 256²), and the
+same species in the source data:
+
+<img src="docs/img/samples.png" width="100%"/>
+
+## Usage
+
+```bash
+python3 src/sample.py --prompt "butterfly, Troides hypolitus" --out out.png
+```
+
+## Training / fine-tuning
+
+```bash
+bash run.sh /path/to/tfrecord_shards 10      # 10 эпох: TFRecord-шарды i1 (PNG 256² + подписи)
+python3 src/train.py                         # дефолты: --init transformer, --out tmp/run, --target 708000
+```
+
+## Limitations
+
+- Butterflies only, right now: 966 training images, 3-word captions.
+- Colour drifts towards saturated species stereotypes; vein pattern is approximate, not faithful.
+- No text rendering, no faces, no general scenes. Research artifact — no safety tuning.
+
+MIT for our code and weights (`src/`, `transformer/`). The bundled VAE is `AsymmetricAutoencoderKL`
+(32 latent channels, encoder f8 / decoder f16); the text encoder is MIT.

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 обнулён, поэтому новые ветки не ломают старт.
+
+Размеры: тело **240.0M** (40 блоков), миксер 10.0M, VAE отдельно (асимметричный: энкодер f8 / декодер f16). Шаг обучения 97 мс при bs8 на 256² (8.1 ГБ). Для сравнения, прежний
+плоский стек ширины 1024 на всех уровнях: 196 мс и hf 0.18 против 0.54 у пирамиды (код и чекпойнт
+той архитектуры удалены как промежуточный артефакт).
+
+## 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. Бюджет
+
+| часть | параметры |
+|---|---|
+| тело | 240.0M (40 блоков) |
+| выбор слоёв текста | 10.04M |
+| VAE | отдельно, `vae/` (AsymmetricAutoencoderKL) |
+| **итого** | **177.3M** |
+
+## 6. Что зафиксировано и почему
+
+- **x0, а не velocity.** Смена параметризации тела не затрагивает декодер (это отдельный flow со
+  своей целью), но требует переучивания тела с нуля — поэтому оставлено как есть.
+- **Текст не экономим.** Кросс-внимание на L0 добавило 2.5M параметров и 9 мс, `switch` вырос
+  0.43 → 0.78 при том же уровне деталей.
+- **Никаких фиксированных шумов и клипов в обучении** — только стандартный flow-matching.
+- **Удалено как тупиковое:** плоский 1024-стек, резидуальная U-ветка, Mix-FFN поверх плоского
+  стека, velocity-переключение посреди обучения, LPIPS в лоссе, кропы одной картинки вместо набора.

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>

requirements.txt

@@ -0,0 +1,27 @@
+# Ra — что нужно, чтобы обучать и сэмплить.
+# Версии не прибиты: нижние границы — то, с чем реально работает код.
+# (проверено на torch 2.6 / transformers 5.5 / diffusers 0.41 / hf_hub 1.33)
+
+# ядро: модель, обучение, сэмплинг
+torch>=2.4
+numpy>=1.24
+Pillow>=10.0
+safetensors>=0.4
+huggingface_hub>=0.27          # hf download / HfFileSystem (поток шардов и датасетов)
+transformers>=4.48             # ModernBertModel (текстовый энкодер Ettin)
+
+# поток данных FLUX-Reason-6M (parquet)
+datasets>=2.19
+pyarrow>=14.0
+
+# 8-битный AdamW; без него код автоматически перейдёт на обычный AdamW
+bitsandbytes>=0.43
+
+# необязательно: triton ускоряет оконное внимание внутри нашего VAE (vae/ae/window_attention.py).
+# Без него работает путь на SDPA — тише, но немного медленнее. Обычно приходит вместе с torch.
+# triton>=2.1
+
+# ВАЖНО: сторонний VAE sdxs (vae_sdxs/, AsymmetricAutoencoderKL) требует diffusers >= 0.37.
+# Наш собственный VAE (vae/, f8c32) идёт своим кодом в vae/ae/ и diffusers не требует.
+# Если ставишь diffusers отдельно и не нужен sdxs — эту строку можно убрать.
+diffusers>=0.37

src/export.py

@@ -0,0 +1,105 @@
+"""Экспорт обученного UDM в `transformer/`: safetensors + config.json.
+
+Три части в одном файле, по префиксам:
+  body.*        — пирамидальное тело, поток в латентном пространстве VAE
+  text_mixer.*  — обучаемый выбор слоёв Ettin (SliceMixer)
+
+  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 = {}
+    body_sd = ck.get("body_ema") or ck["body"]        # для публикации берём EMA, если обучен
+    mixer_sd = ck.get("mixer_ema") or ck["mixer"]
+    for name, sd in (("body", body_sd), ("text_mixer", mixer_sd)):
+        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())   # AsymmetricAutoencoderKL
+    ta = ck.get("args") or {}
+    counts = {}                          # сколько блоков в каждом списке — по ключам тензоров
+    for k in ck["body"]:
+        g = k.split(".")[0]
+        if g[0] in "ed" and "." in k:
+            idx = k.split(".")[1]
+            if idx.isdigit():
+                counts.setdefault(g, set()).add(idx)
+    nblk = {g: len(v) for g, v in counts.items()}
+    # конфиг описывает ТОЛЬКО тело; статистики латента — в vae/config.json, шаг — параметр запуска
+    cfg = {
+        "_class_name": "UDiT",
+        "format": "udm/v1",
+        "storage_dtype": "float16" if args.half else "float32",
+        "latent_channels": ae_cfg["latent_channels"],
+        "text_dim": ta.get("text_dim", 768),
+        "body": {
+            "kind": "пирамидальный U: ширина растёт к низу",
+            "widths": [192, 512, 1024, 1536],
+            "heads": [6, 8, 8, 8],
+            "blocks": nblk,
+            "levels": [
+                {"level": 0, "tokens": "÷1", "width": 192,
+                 "blocks": "conv 6 enc / 4 dec + 2 self-attn на сторону"},
+                {"level": 1, "tokens": "÷2", "width": 512, "blocks": "attn 5 enc / 4 dec"},
+                {"level": 2, "tokens": "÷4", "width": 1024, "blocks": "attn 5 enc / 4 dec"},
+                {"level": 3, "tokens": "÷8", "width": 1536, "blocks": "attn 4 (боттлнек)"},
+            ],
+            "skips": "аддитивные (SD-стиль): e0/e1/e2 прибавляются в декодере",
+            "down": "mean-pool 2x + Linear",
+            "up": "Linear + nearest 2x",
+            "positions": {
+                "rope": {"kind": "2d", "base": 100, "phase": "в латентных пикселях на каждом уровне"},
+                "absolute": {"kind": "2d sincos", "dim": 64,
+                             "into": "zero-init Linear(64 -> width) к выходу stem, в паре с RoPE"},
+            },
+            "text": "кросс-внимание на всех уровнях",
+            "parametrization": "x0 латента",
+            "params": round(sum(v.numel() for k, v in tensors.items() if k.startswith("body.")) / 1e6, 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['body']['params']}M, блоки {nblk}")
+
+
+def load_exported(model, enc, ae, path, device=None):
+    """Загрузить из `transformer/` (config.json + safetensors с префиксами body./text_mixer.)."""
+    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"))
+    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,344 @@
+"""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"))
+
+
+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)))
+
+
+def zero_out(blk):
+    """Обнулить выходы внимания: новый блок на старте тождественен (FFN-выход и так zero-init)."""
+    for name in ("o", "xo"):
+        lin = getattr(blk, name, None)
+        if isinstance(lin, nn.Linear):
+            nn.init.zeros_(lin.weight)
+            if lin.bias is not None:
+                nn.init.zeros_(lin.bias)
+    return blk
+
+
+GRAFT_DEFAULT = {"e0x": 3, "d0x": 3, "e1": 1, "d1": 1, "e2": 2, "d2": 2, "e3": 2}
+
+
+def sincos_2d(gh: int, gw: int, dim: int, device, dtype=None):
+    """Абсолютные 2D-позиции (DiT-стиль sin/cos по каждой оси) — в паре с RoPE в attention."""
+    key = ("pe2d", gh, gw, dim, str(device))
+    tab = _ROPE_CACHE.get(key)
+    if tab is None:
+        half = dim // 2
+        f = 1.0 / (10000 ** (torch.arange(half // 2, device=device, dtype=torch.float32) / (half / 2)))
+        y = torch.arange(gh, device=device, dtype=torch.float32)[:, None] * f[None]
+        x = torch.arange(gw, device=device, dtype=torch.float32)[:, None] * f[None]
+        ey = torch.cat([y.sin(), y.cos()], -1)[:, None, :].expand(gh, gw, half)
+        ex = torch.cat([x.sin(), x.cos()], -1)[None, :, :].expand(gh, gw, half)
+        tab = torch.cat([ey, ex], -1).reshape(gh * gw, dim)
+        _ROPE_CACHE[key] = tab
+    return tab if dtype is None else tab.to(dtype)
+
+
+class Body(nn.Module):
+    """U-стек: ширина 192 → 512 → 1024 → 1536 с аддитивными скипами.
+
+    Базовая архитектура = 36 блоков: L0 10 conv + 4 attn, L1 9, L2 9, L3 4 (238.5M, 44.4 GFLOP на 256 px).
+    Привитые блоки входят в GRAFT_DEFAULT и на конструкторе тождественны — так старые чекпойнты
+    (26 блоков) грузятся без сломки: лишние блоки остаются тождественными.
+    """
+
+    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
+        self.tdim = text_dim if xatt else None
+        heads = {0: 6, 1: 8, 2: 8, 3: 8}
+        self.stem = nn.Linear(latent_ch, w0)
+        self.pe_dim = 64
+        self.pe = nn.Linear(self.pe_dim, w0, bias=False)      # zero-init: на старте не вмешивается
+        nn.init.zeros_(self.pe.weight)
+        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.e0x = nn.ModuleList()     # привива��мые L0-attn (пусто = архитектура без изменений)
+        self.d0x = nn.ModuleList()
+        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)
+        self.graft(GRAFT_DEFAULT)     # базовая архитектура уже с прививкой (блоки тождественны)
+
+    # ---- прививка блоков
+    def graft(self, plan: dict) -> int:
+        """Добавить блоки к уже обученной модели так, чтобы они стартовали тождеством.
+
+        plan: {"e1": 2, "d1": 2, "e3": 2}; ключи — имена списков блоков.
+        Тождество: FFN-выход (`down`) zero-init, выходы внимания (`o`, `xo`) обнуляем здесь.
+        """
+        heads = {0: 6, 1: 8, 2: 8, 3: 8}
+        lvl = {"e0": 0, "d0": 0, "e0x": 0, "d0x": 0, "e1": 1, "d1": 1, "e2": 2, "d2": 2, "e3": 3}
+        dev = next(self.parameters()).device
+        added = 0
+        for key, n in plan.items():
+            seq = getattr(self, key)
+            for _ in range(int(n)):
+                if key in ("e0", "d0"):
+                    blk = ConvBlock(self.widths[0], text_dim=self.tdim)
+                else:
+                    blk = AttnBlock(self.widths[lvl[key]], heads[lvl[key]], self.tdim)
+                seq.append(zero_out(blk).to(dev))
+                added += 1
+        return added
+
+    # ---- переходы
+    @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
+        # pooled-текст в adaLN: среднее ТОЛЬКО по реальным токенам (паддинг иначе разбавляет вектор,
+        # и поведение модели начинает зависеть от длины паддинга, а не от подписи)
+        # pooled-текст для adaLN: среднее ТОЛЬКО по реальным токенам. Иначе вектор зависит от
+        # длины паддинга, а не от подписи (обучение паддит до длины подписи, инференс — до 192)
+        m = text_mask[..., None].to(text.dtype)
+        pooled = (text * m).sum(1) / m.sum(1).clamp(min=1.0)
+        cond = self.tmlp(t_embed(t, 256)) + self.pmlp(pooled)
+        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))
+        x = x + self.pe(sincos_2d(gh0, gw0, self.pe_dim, x.device, x.dtype))
+        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)
+        for blk in self.e0x:
+            x = blk(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.d0x:
+            x = blk(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, 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,113 @@
+#!/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", "requirements.txt", "run_sdxs.sh", ".gitattributes", "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",
+        ]
+
+
+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("--readme", default="README.md", help="какой README положить в корень")
+    ap.add_argument("--push", action="store_true", help="залить собранное на Hub (публичный репозиторий)")
+    ap.add_argument("--step", type=int, default=None, help="шаг чекпойнта — только для сообщения коммита")
+    args = ap.parse_args()
+    out = Path(args.out)
+    if out.exists():
+        shutil.rmtree(out)
+    out.mkdir(parents=True)
+
+    rels = list(COPY) + [args.readme, "vae/config.json", "transformer/config.json"]
+    for rel in rels:
+        src = ROOT / rel
+        if not src.exists():
+            continue
+        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:
+            dst_rel = "README.md" if rel == args.readme else rel
+            shutil.copy2(src, (out / dst_rel))
+
+    tdir, vdir = ROOT / "transformer", ROOT / "vae"
+    half_file(tdir / "diffusion_pytorch_model.safetensors",
+              out / "transformer" / "diffusion_pytorch_model.safetensors")
+    vae_weights = vdir / "model.safetensors"
+    if not vae_weights.exists():                       # AsymmetricAutoencoderKL кладёт так
+        vae_weights = vdir / "diffusion_pytorch_model.safetensors"
+    half_file(vae_weights, out / "vae" / vae_weights.name)
+
+    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 CommitOperationAdd, HfApi
+        api = HfApi()
+        api.create_repo(args.repo, repo_type="model", private=False, exist_ok=True)
+        # create_commit вместо upload_folder: upload_folder уважает .gitignore из бандла
+        # (а там шаблоны `*.safetensors`) и молча не зали��ает веса
+        ops = [CommitOperationAdd(path_in_repo=p.relative_to(out).as_posix(), path_or_fileobj=str(p))
+               for p in sorted(out.rglob("*")) if p.is_file()]
+        api.create_commit(repo_id=args.repo, repo_type="model", operations=ops,
+                          commit_message=(f"Ra sdxs-VAE (step {args.step})" if args.step
+                                          else "Ra sdxs-VAE (weights)"))
+        print(f"[push] https://huggingface.co/{args.repo}: {len(ops)} файлов", flush=True)
+
+
+if __name__ == "__main__":
+    main()

src/sample.py

@@ -0,0 +1,74 @@
+#!/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 vae_adapter import VAE          # 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("--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 = VAE.from_pretrained(str(ROOT / "vae"), 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}" + (f", шаг {step}" if step else ""), 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=0, 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/tfrecord_stream.py

@@ -0,0 +1,208 @@
+"""Поток картинок из TFRecord-шардов i1 (PNG 256×256 + 5 подписей) вместо parquet-потока FLUX.
+
+Отличие от HFImageStream: запись TFRecord = [uint64 len][crc][payload][crc], внутри PNG-байты
+и `caption` — пять длинных перефразировок одной картинки. Поэтому:
+  * из пяти подписей берём случайно самую короткую или следующую по длине (бюджета на все нет,
+    а короткая лучше влезает в --text-len);
+  * картинки уже 256×256 — ресайз только если сторона не совпала с --size;
+  * источник: локальная папка с шардами (быстро) или прямо репозиторий HF через range-запросы
+    (block_size 16 МБ даёт ~90 записей/с, ~9.7 МБ/с — на уровне нашего темпа обучения).
+"""
+from __future__ import annotations
+
+import io
+import queue
+import random
+import struct
+import threading
+from pathlib import Path
+
+import numpy as np
+import torch
+from PIL import Image
+
+SHARD_GLOB = "*.tfrecord-*"
+
+
+# ---------------------------------------------------------------- разбор формата
+def _varint(buf: bytes, i: int) -> tuple[int, int]:
+    shift = val = 0
+    while True:
+        b = buf[i]; i += 1
+        val |= (b & 0x7F) << shift
+        if not b & 0x80:
+            return val, i
+        shift += 7
+
+
+def _fields(buf: bytes):
+    i = 0
+    while i < len(buf):
+        key, i = _varint(buf, i)
+        num, wire = key >> 3, key & 7
+        if wire == 2:
+            ln, i = _varint(buf, i)
+            yield num, buf[i:i + ln]
+            i += ln
+        elif wire == 0:
+            v, i = _varint(buf, i)
+            yield num, v
+        elif wire == 1:
+            yield num, buf[i:i + 8]; i += 8
+        elif wire == 5:
+            yield num, buf[i:i + 4]; i += 4
+        else:
+            raise ValueError(f"wire type {wire}")
+
+
+def _packed_i64(buf: bytes) -> list[int]:
+    out, i = [], 0
+    while i < len(buf):
+        v, i = _varint(buf, i)
+        out.append(v)
+    return out
+
+
+def parse_example(payload: bytes) -> dict:
+    """tf.train.Example → {поле: список bytes | список int}"""
+    out = {}
+    for num, features in _fields(payload):
+        if num != 1:
+            continue
+        for fn, entry in _fields(features):
+            if fn != 1:
+                continue
+            key, val = None, {}
+            for en, ev in _fields(entry):
+                if en == 1:
+                    key = ev.decode()
+                elif en == 2:
+                    for fnum, fval in _fields(ev):
+                        if fnum == 1:
+                            val["bytes"] = [v for n, v in _fields(fval) if n == 1]
+                        elif fnum == 3:
+                            val["int64"] = [x for n, v in _fields(fval) if n == 1 for x in _packed_i64(v)]
+            out[key] = val.get("bytes", val.get("int64"))
+    return out
+
+
+def count_records(paths) -> int:
+    """Сколько записей в шардах — читаем только заголовки (seek через payload), быстро."""
+    total = 0
+    for path in paths:
+        with open(path, "rb") as fh:
+            size = fh.seek(0, 2)
+            pos = 0
+            while pos + 12 <= size:
+                fh.seek(pos)
+                head = fh.read(8)
+                if len(head) < 8:
+                    break
+                (n,) = struct.unpack("<Q", head)
+                pos += 8 + 4 + n + 4
+                total += 1
+    return total
+
+
+def count_shard_records(source: str) -> int | None:
+    """Точное число записей для локальной папки; None, если источник удалённый."""
+    if not Path(source).is_dir():
+        return None
+    paths = sorted(p.as_posix() for p in Path(source).glob(SHARD_GLOB))
+    return count_records(paths) if paths else None
+
+
+def read_records(fh):
+    """файловый объект (локальный или HfFileSystem) → поток payload'ов"""
+    while True:
+        head = fh.read(8)
+        if len(head) < 8:
+            return
+        (n,) = struct.unpack("<Q", head)
+        fh.read(4)
+        payload = fh.read(n)
+        fh.read(4)
+        yield payload
+
+
+# ---------------------------------------------------------------- поток
+class TFRecordStream:
+    def __init__(self, source: str, size: int = 256, workers: int = 8, prefetch: int = 256,
+                 block_mb: int = 16, seed: int = 0, caption_short2: bool = True):
+        """source — локальная папка с шардами или id датасета на HF (repo_type=dataset)."""
+        self.q: queue.Queue = queue.Queue(maxsize=prefetch)
+        self.stop = False
+        self.size, self.workers, self.seed = size, workers, seed
+        self.caption_short2 = caption_short2
+        self.local = Path(source).is_dir()
+        self.source = str(source)
+        self.block = block_mb * 2 ** 20
+        threading.Thread(target=self._loop, daemon=True).start()
+
+    # ---- список шардов
+    def _shards(self) -> list[str]:
+        if self.local:
+            files = sorted(p.as_posix() for p in Path(self.source).glob(SHARD_GLOB))
+        else:
+            from huggingface_hub import HfFileSystem
+            fs = HfFileSystem()
+            files = [f"datasets/{self.source}/{p}" for p in sorted(fs.ls(f"datasets/{self.source}"))
+                     if ".tfrecord-" in p]
+        if not files:
+            raise FileNotFoundError(f"нет шардов в {self.source}")
+        random.Random(self.seed).shuffle(files)
+        return files
+
+    def _open(self, path: str):
+        if self.local:
+            return open(path, "rb")
+        from huggingface_hub import HfFileSystem
+        return HfFileSystem().open(path, "rb", block_size=self.block)
+
+    # ---- подготовка одной записи
+    def _prepare(self, payload: bytes):
+        d = parse_example(payload)
+        im = Image.open(io.BytesIO(d["image"][0])).convert("RGB")
+        w, h = im.size
+        if min(w, h) != self.size:
+            sc = self.size / min(w, h)
+            im = im.resize((max(self.size, round(w * sc)), max(self.size, round(h * sc))), Image.LANCZOS)
+            w, h = im.size
+            im = im.crop(((w - self.size) // 2, (h - self.size) // 2,
+                          (w - self.size) // 2 + self.size, (h - self.size) // 2 + self.size))
+        caps = [c.decode() if isinstance(c, bytes) else str(c) for c in d.get("caption", [])]
+        caps = sorted(caps, key=len)
+        if not caps:
+            caps = [""]
+        cap = random.choice(caps[:2]) if self.caption_short2 else random.choice(caps)
+        arr = torch.from_numpy(np.asarray(im, dtype=np.float32)).permute(2, 0, 1) / 127.5 - 1.0
+        return arr, cap
+
+    def _loop(self):
+        import collections
+        from concurrent.futures import ThreadPoolExecutor
+        with ThreadPoolExecutor(max_workers=self.workers) as pool:
+            while not self.stop:
+                try:
+                    for shard in self._shards():
+                        fh = self._open(shard)
+                        try:
+                            it = read_records(fh)
+                            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())
+                        finally:
+                            fh.close()
+                except Exception as e:
+                    print(f"[tfrecord] сбой ({type(e).__name__}: {e}), переподключение через 5 с", flush=True)
+                    import time
+                    time.sleep(5)
+
+    def batch(self, n: int):
+        imgs, caps = [], []
+        while len(imgs) < n:
+            arr, cap = self.q.get()
+            imgs.append(arr); caps.append(cap)
+        return torch.stack(imgs), caps

src/train.py

@@ -0,0 +1,604 @@
+#!/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), switch, 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 vae_adapter import VAE                  # noqa: E402
+from tfrecord_stream import TFRecordStream, count_shard_records   # noqa: E402  (шарды i1)
+
+
+def load_ae(args):
+    """VAE проекта: асимметричный, декодер отдаёт 2x. Хвоста и пиксельного лосса нет."""
+    return VAE.from_pretrained(str(ROOT / "vae"), DEV)
+
+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=256,
+                 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("--tfrecord", default="",
+                    help="папка с TFRecord-шардами i1 или id датасета на HF; приоритетнее --hf-dataset")
+    ap.add_argument("--hf-dataset", default="LucasFang/FLUX-Reason-6M",
+                    help="поток HF; пустая строка = папка --data")
+    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=256)
+    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=None, help="бюджет в шагах (по умолчанию — до --target)")
+    ap.add_argument("--target", type=int, default=708000,
+                    help="абсолютный шаг финиша, если не заданы --steps/--epochs")
+    ap.add_argument("--step", type=int, default=0,
+                    help="стартовый шаг: от него считается остаток и расписание lr")
+    ap.add_argument("--batch", type=int, default=16)
+    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=4.5e-5)
+    ap.add_argument("--lr-warmup", type=int, default=0, help="шагов прогрева (0 — продолжаем с готовых весов)")
+    ap.add_argument("--lr-min", type=float, default=1e-5, help="нижняя точка cosine")
+    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("--graft", default="", help="привить блоки (стартуют тождеством), например 'e0=2,e1=2,d1=2'")
+    ap.add_argument("--cfg", type=float, default=2.5, help="CFG при показе сэмплов")
+    ap.add_argument("--sample-steps", type=int, default=16, help="шагов тела при сэмплинге")
+    ap.add_argument("--init", default=str(ROOT / "transformer"),
+                    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=2000, help="0 = не сэмплировать")
+    ap.add_argument("--save-every", type=int, default=4000, 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("--ema", type=float, default=0.999, help="EMA весов для сэмплинга (0 = выкл)")
+    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,img_s,eta_h,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.tfrecord:
+        stream = TFRecordStream(args.tfrecord, size=args.size, workers=args.workers,
+                                prefetch=args.prefetch, seed=seed)
+        counted = count_shard_records(args.tfrecord)
+        n = counted or 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 or 0)
+        print(f"[data] TFRecord {args.tfrecord}: "
+              + (f"{n} картинок в папке" if counted else f"~{n} картинок (оценка)")
+              + f", эпоха = {per_epoch} шагов, подпись = случайная из двух самых коротких → "
+              + (f"бюджет {steps} шагов (~{steps * args.batch / 140 / 3600:.1f} ч при 140 карт/с)"
+                 if steps else f"цель шаг {args.target}"), flush=True)
+        ae = load_ae(args)
+        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)
+    elif 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 or 0)
+        print(f"[data] поток {args.hf_dataset} [{args.parts}], ~{n} картинок/эпоха, "
+              f"{per_epoch} шагов/эпоха → "
+              + (f"бюджет {steps} шагов (~{steps * args.batch / 62 / 3600:.1f} ч при 62 картинках/с)"
+                 if steps else f"цель шаг {args.target}"), flush=True)
+        ae = load_ae(args)
+        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 = load_ae(args)
+        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 or 0)
+        print(f"[data] {n} пар, латент {tuple(Z.shape[1:])}, текст {tuple(HS.shape[1:])}, "
+              f"эпоха = {per_epoch} шагов → "
+              + (f"бюджет {steps}" if steps else f"цель шаг {args.target}"), 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.graft:
+        plan = {k.strip(): int(v) for k, v in (kv.split("=") for kv in args.graft.split(","))}
+        added = m.body.graft(plan)
+        print(f"[graft] привито {added} блоков: {plan}", flush=True)
+
+    if args.init:
+        p = Path(args.init)
+        if p.is_dir():
+            from export import load_exported
+            load_exported(m, enc, ae, p, DEV)
+            start = args.step                    # шаг — параметр запуска, не поле конфига
+            print(f"[init] папка {p}; старт шага {start} (--step)", flush=True)
+        else:
+            sd = torch.load(p, map_location=DEV)
+            n_take = load_body(sd["body"]); enc.mixer.load_state_dict(sd["mixer"])
+            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"])
+        start = int(sd.get("step", 0))
+        print(f"[resume] {out / 'last.pt'} → шаг {start}", flush=True)
+
+    if not args.steps and not args.epochs:            # бюджет по умолчанию — до абсолютного шага
+        steps = max(0, args.target - start)
+        print(f"[budget] цель шаг {args.target} от {start} → {steps} шагов", flush=True)
+    if steps <= 0:
+        print(f"[budget] цель {args.target} уже достигнута (шаг {start}) — учить нечего", flush=True)
+        if stream is not None:
+            stream.stop = True
+        import os as _os
+        _os._exit(0)
+
+    n_body = sum(p.numel() for p in m.body.parameters())
+    print(f"[model] тело {n_body/1e6:.1f}M + миксер {sum(p.numel() for p in enc.trainable())/1e6:.2f}M",
+          flush=True)
+
+    groups = [{"params": list(m.body.parameters()), "lr": args.lr},
+              {"params": enc.trainable(), "lr": args.lr}]
+    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)
+
+    # ---- EMA весов: сглаженные копии для сэмплинга (обучение идёт на сырых)
+    ema_body = ema_mixer = None
+    if args.ema:
+        ema_body = {k: p.detach().float().clone() for k, p in m.body.named_parameters()}
+        ema_mixer = {k: p.detach().float().clone() for k, p in enc.mixer.named_parameters()}
+        print(f"[ema] decay {args.ema}, {sum(v.numel() for v in ema_body.values())/1e6:.1f}M параметров тела", flush=True)
+
+    def ema_update():
+        if ema_body is None:
+            return
+        with torch.no_grad():
+            for k, p in m.body.named_parameters():
+                ema_body[k].mul_(args.ema).add_(p.detach().float(), alpha=1 - args.ema)
+            for k, p in enc.mixer.named_parameters():
+                ema_mixer[k].mul_(args.ema).add_(p.detach().float(), alpha=1 - args.ema)
+
+    class _EmaSwap:
+        """Подменить веса на EMA на время сэмплинга, потом вернуть сырые."""
+        def __enter__(self):
+            if ema_body is None:
+                return self
+            self.raw = {k: p.detach().cpu().clone() for k, p in m.body.named_parameters()}
+            self.raw_mx = {k: p.detach().cpu().clone() for k, p in enc.mixer.named_parameters()}
+            with torch.no_grad():
+                for k, p in m.body.named_parameters():
+                    p.copy_(ema_body[k].to(p.device))
+                for k, p in enc.mixer.named_parameters():
+                    p.copy_(ema_mixer[k].to(p.device))
+            return self
+        def __exit__(self, *exc):
+            if ema_body is None:
+                return False
+            with torch.no_grad():
+                for k, p in m.body.named_parameters():
+                    p.copy_(self.raw[k].to(p.device))
+                for k, p in enc.mixer.named_parameters():
+                    p.copy_(self.raw_mx[k].to(p.device))
+            return False
+
+    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_slices, show_x = [], []          # срезы замороженного энкодера; миксер учится — смешиваем при каждом показе
+    if stream is not None:
+        for cap in SAMPLE_PROMPTS:
+            show_slices.append(enc.slices([cap]))
+    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]]
+        if len(show) < 4:            # один класс в eval (бабочки) — берём первые 4 файла как есть
+            show = 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_slices.append((hs, mk))
+            show_x.append(load(f)[None].to(DEV))
+
+    def sample_grid(step):
+        tiles, hfs, ys = [], [], []
+        # веса уже подменены на EMA в вызывающем месте
+        for hs, mk in show_slices:
+            y = m.sample(enc.mix(hs.float()), mk, lat_hw[0], lat_hw[1], steps=args.sample_steps,
+                         tail_steps=0, 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)
+            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()},
+              "body_ema": ({k: v.cpu() for k, v in ema_body.items()} if ema_body else None),
+              "mixer_ema": ({k: v.cpu() for k, v in ema_mixer.items()} if ema_mixer else None),
+              "mixer": {k: v.cpu() for k, v in enc.mixer.state_dict().items()},
+              "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)
+    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()
+            total = loss_l
+            if not torch.isfinite(total):
+                skips += 1
+                continue
+            (total / args.accum).backward()
+            ok += 1
+            run["lat"] += loss_l.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)):
+            g["lr"] = lr * (base / args.lr)
+        opt.step()
+        ema_update()
+
+        if step % args.log_every == 0:
+            d = args.ema_decay
+            cur = {"lat": run["lat"] / 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
+            ms = inc / max(run["n"], 1) * 1e3
+            avg = el / max(step - start, 1)                 # с/шаг с начала прогона (включая простои)
+            left = (start + steps - step) * avg
+            eta = time.strftime("%d.%m %H:%M", time.gmtime(time.time() + left))
+            print(f"  {step:6d} | mse_lat {cur['lat']:.4f} (ema {ema['lat']:.4f})"
+                  f" | lr {lr:.1e} | {args.batch * args.accum / (ms / 1e3):.1f} карт/с ({ms:.0f} мс/шаг)"
+                  f" | осталось {left / 3600:.1f} ч (финиш {eta} UTC)"
+                  + (f" | skip {skips}" if skips else ""), 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:
+            ema_ctx = _EmaSwap().__enter__()
+            if DEV == "cuda":
+                torch.cuda.empty_cache()      # сэмплинг идёт поверх фрагментированного кэша — чистим заранее
+            zs = []
+            lat_pc_min = float("nan")
+            for hs, mk in show_slices[:2]:
+                txt = enc.mix(hs.float())
+                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)
+            tail_psnr = float("nan")          # метрика хвоста больше не существует
+            el = time.perf_counter() - t0
+            avg = el / max(step - start, 1); left = (start + steps - step) * avg
+            print(f"  ↳ {step} | hf {hfg:.3f} | латент std {lat_pc_min:.2f} | switch {switch:.3f}"
+                  f" | {time.strftime('%H:%M', time.gmtime())} UTC", 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"0.0000,{switch:.4f},{bright:.4f},{contrast:.4f},"
+                         f"{lat_pc_min:.4f},{hfg:.4f},{tail_psnr:.2f},"
+                         f"{args.batch * args.accum / (el / max(step - start, 1)):.1f},{left / 3600:.2f},{el:.1f}\n")
+            ema_ctx.__exit__(None, None, None)      # вернуть сырые веса
+        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()

src/vae_adapter.py

@@ -0,0 +1,61 @@
+"""VAE проекта: AsymmetricAutoencoderKL (энкодер f8, декодер f16 → выход 2×).
+
+Что адаптер прячет от остального кода:
+  * латент нормируется своими `latents_mean`/`latents_std` из vae/config.json (≈ -0.04 / 1.58):
+    encode(x) → нормированный латент, `to_raw()` в модели возвращает сырой;
+  * декодер асимметричный: энкодер f8, декодер f16, то есть латент 32×32 отдаёт картинку 512²
+    (вдвое больше входа) — это и есть смысл этого VAE;
+  * декодирование однопроходное — обучение идёт как обычная латентная диффузия (только MSE по латенту).
+
+API, которого ждут train.py/sample.py: `encode`, `decode_raw`, `latent_mean`, `latent_std`,
+`config.latent_channels`, `config.patch_size`, `compute_dtype`, `device`, `.eval()`, `.to(device)`.
+"""
+from __future__ import annotations
+
+import torch
+import torch.nn as nn
+
+PATCH = 8                      # энкодер сжимает в 8 раз; декодер потом растягивает в 16 → 2× на выходе
+
+
+class VAE(nn.Module):
+    def __init__(self, vae, device: str = "cuda"):
+        super().__init__()
+        self.vae = vae
+        self.device_name = device
+        self.compute_dtype = torch.float32
+        cfg = vae.config
+        mean = torch.tensor(getattr(cfg, "latents_mean"), dtype=torch.float32)
+        std = torch.tensor(getattr(cfg, "latents_std"), dtype=torch.float32)
+        self.register_buffer("latent_mean", mean.view(1, -1, 1, 1))
+        self.register_buffer("latent_std", std.view(1, -1, 1, 1))
+        self.config = type("Cfg", (), {"latent_channels": int(cfg.latent_channels),
+                                       "patch_size": PATCH})()
+
+    @classmethod
+    def from_pretrained(cls, path, device: str = "cuda"):
+        from diffusers import AsymmetricAutoencoderKL
+        vae = AsymmetricAutoencoderKL.from_pretrained(str(path), torch_dtype=torch.float32)
+        vae = vae.to(device).eval()
+        vae.requires_grad_(False)
+        return cls(vae, device).to(device).eval()
+
+    @property
+    def device(self):
+        return self.latent_mean.device
+
+    def train(self, mode: bool = True):
+        super().train(mode)
+        self.vae.eval()                     # VAE заморожен всегда
+        return self
+
+    def encode(self, x: torch.Tensor) -> torch.Tensor:
+        """x в -1..1 → нормированный латент (B, 32, H/8, W/8)."""
+        with torch.no_grad():
+            raw = self.vae.encode(x.to(self.compute_dtype)).latent_dist.mode()
+        return ((raw - self.latent_mean) / self.latent_std).to(torch.float32)
+
+    def decode_raw(self, z_raw: torch.Tensor, steps: int = 0, generator=None) -> torch.Tensor:
+        """Сырой латент → картинка в -1..1. `steps`/`generator` — для совместимости с нашим API."""
+        with torch.no_grad():
+            return self.vae.decode(z_raw.to(self.compute_dtype), return_dict=False)[0].clamp(-1, 1)

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,77 @@
+{
+  "_class_name": "UDiT",
+  "format": "udm/v1",
+  "storage_dtype": "float16",
+  "latent_channels": 32,
+  "text_dim": 768,
+  "body": {
+    "kind": "пирамидальный U: ширина растёт к низу",
+    "widths": [
+      192,
+      512,
+      1024,
+      1536
+    ],
+    "heads": [
+      6,
+      8,
+      8,
+      8
+    ],
+    "blocks": {
+      "e0": 6,
+      "e0x": 3,
+      "d0x": 3,
+      "e1": 5,
+      "e2": 5,
+      "e3": 4,
+      "d2": 4,
+      "d1": 4,
+      "d0": 4
+    },
+    "levels": [
+      {
+        "level": 0,
+        "tokens": "÷1",
+        "width": 192,
+        "blocks": "conv 6 enc / 4 dec + 2 self-attn на сторону"
+      },
+      {
+        "level": 1,
+        "tokens": "÷2",
+        "width": 512,
+        "blocks": "attn 5 enc / 4 dec"
+      },
+      {
+        "level": 2,
+        "tokens": "÷4",
+        "width": 1024,
+        "blocks": "attn 5 enc / 4 dec"
+      },
+      {
+        "level": 3,
+        "tokens": "÷8",
+        "width": 1536,
+        "blocks": "attn 4 (боттлнек)"
+      }
+    ],
+    "skips": "аддитивные (SD-стиль): e0/e1/e2 прибавляются в декодере",
+    "down": "mean-pool 2x + Linear",
+    "up": "Linear + nearest 2x",
+    "positions": {
+      "rope": {
+        "kind": "2d",
+        "base": 100,
+        "phase": "в латентных пикселях на каждом уровне"
+      },
+      "absolute": {
+        "kind": "2d sincos",
+        "dim": 64,
+        "into": "zero-init Linear(64 -> width) к выходу stem, в паре с RoPE"
+      }
+    },
+    "text": "кросс-внимание на всех уровнях",
+    "parametrization": "x0 латента",
+    "params": 240.04
+  }
+}

transformer/diffusion_pytorch_model.safetensors

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

vae/config.json

@@ -0,0 +1,108 @@
+{
+  "_class_name": "AsymmetricAutoencoderKL",
+  "_diffusers_version": "0.37.0",
+  "_name_or_path": "asymmetric_vae",
+  "act_fn": "silu",
+  "down_block_out_channels": [
+    128,
+    256,
+    512,
+    512
+  ],
+  "down_block_types": [
+    "DownEncoderBlock2D",
+    "DownEncoderBlock2D",
+    "DownEncoderBlock2D",
+    "DownEncoderBlock2D"
+  ],
+  "in_channels": 3,
+  "latent_channels": 32,
+  "latents_mean": [
+    -0.03542253375053406,
+    0.20086465775966644,
+    -0.016413161531090736,
+    -0.0956302210688591,
+    -0.2672063112258911,
+    0.2609933018684387,
+    -0.07806991040706635,
+    -0.48407721519470215,
+    0.21844269335269928,
+    -0.1122383326292038,
+    0.27197545766830444,
+    -0.18958772718906403,
+    0.18776826560497284,
+    0.0987580344080925,
+    0.2837068736553192,
+    -0.4486690163612366,
+    0.4816776514053345,
+    0.02947971224784851,
+    -0.1337375044822693,
+    -0.39750921726226807,
+    -0.08513020724058151,
+    -0.054023586213588715,
+    -0.3943594992160797,
+    0.23918119072914124,
+    -0.12466679513454437,
+    0.09935147315263748,
+    0.31858691573143005,
+    0.48585832118988037,
+    -0.6416525840759277,
+    -0.15164820849895477,
+    -0.4693508744239807,
+    -0.13071806728839874
+  ],
+  "latents_std": [
+    1.5792087316513062,
+    1.5769503116607666,
+    1.5864241123199463,
+    1.6454921960830688,
+    1.5336694717407227,
+    1.5587652921676636,
+    1.5838669538497925,
+    1.5659377574920654,
+    1.6860467195510864,
+    1.5192310810089111,
+    1.573639988899231,
+    1.5953549146652222,
+    1.5271092653274536,
+    1.6246271133422852,
+    1.7054023742675781,
+    1.607722282409668,
+    1.558642864227295,
+    1.5824549198150635,
+    1.6202995777130127,
+    1.6206320524215698,
+    1.6379750967025757,
+    1.6527063846588135,
+    1.498811960220337,
+    1.5706247091293335,
+    1.5854856967926025,
+    1.4828169345855713,
+    1.5693111419677734,
+    1.692481517791748,
+    1.6409776210784912,
+    1.6216280460357666,
+    1.6087706089019775,
+    1.5776633024215698
+  ],
+  "layers_per_down_block": 2,
+  "layers_per_up_block": 2,
+  "norm_num_groups": 32,
+  "out_channels": 3,
+  "sample_size": 32,
+  "scaling_factor": 1.0,
+  "up_block_out_channels": [
+    128,
+    128,
+    256,
+    512,
+    512
+  ],
+  "up_block_types": [
+    "UpDecoderBlock2D",
+    "UpDecoderBlock2D",
+    "UpDecoderBlock2D",
+    "UpDecoderBlock2D",
+    "UpDecoderBlock2D"
+  ]
+}

vae/diffusion_pytorch_model.safetensors

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