src — чем собран трансформер в этой папке

Здесь сжатие DiT и обучение на своём датасете: 30 блоков вместо 42, 3.723 B вместо 5.082 B (−26.7 %) — качество на уровне исходной модели на прозе, людях и анатомии. Итоговые веса лежат в ../transformer/.

Чтобы повторить прогон, нужен исходный 42-блочный трансформер: он удалён из папки модели, скачать заново из OpenMOSS-Team/MOSS-Image-Base в transformer/ (скрипты читают именно этот путь).

файл	что делает
`window.py`	всё сжатие: `cache` (состояния учителя) → `train` (LoRA окна) → `eval` (замер окна) → `combo` (собрать модель из слияний)
`lab.py`	харнесс: сборка пайплайна, текстовые условия, сэмплинг, LPIPS + CLIP, удаление блоков
`run_hi.sh`	точная последовательность, которой собран трансформер в этой папке
`make_dataset.py`	картинки `.jpg` + парные `.txt` → латенты VAE, условие адаптера и маска; кладёт чанки `chunk_/sub_` в `--save` (+ `empty_condition.npz` — негатив CFG)
`train.py`	файнтюн DiT на этом датасете под 16 ГБ: чекпойнтинг по блокам, Adam по блоку, paged-оптимизатор; читает папку с чанками целиком
`work/`	кэши, чекпойнты LoRA, прогоны и метрики (пересоздаётся, в артефакт не входит)

Методика

Идея. Блоки в середине DiT меняют поток мало и частично отменяют друг друга. Значит из двух-трёх смежных блоков можно оставить один, добавив ему низкоранговую поправку, обученную повторять поведение исходной модели. Глубина падает, переобучение всей сети не нужно.

Схема окна. Блок s остаётся и получает LoRA (rank 32 на q,k,v,o,ffn_up_gate,ffn_down, 1.78 M параметров); блоки s+1, s+2 удаляются. LoRA учится так, чтобы выход блока s совпал с выходом блока s+2 исходной модели.

Точный протокол (window.py cache + window.py train):

Кэш: для N сэмплов (промпт из ~/dataset/civit/*.txt, случайный латент, случайный t) делаем пару грубых шагов Эйлера исходной моделью, чтобы латент был похож на точку реальной траектории, затем один forward с хуками. Пишем вход блока s (выход блока s-1) и выход блока s+2, плюс rope и modulation (без них блок не посчитать).
Обучение: собранный из весов один блок (226 МБ) + LoRA, всё остальное не грузится вообще. Лосс — относительная MSE по токенам ‖pred−target‖²/‖target‖² плюс 0.1·(1−cos), чтобы держать и направление. Батчи собираются по одинаковой длине последовательности (у промптов разное число токенов, паддинг менял бы поведение блока). AdamW8bit, lr 1e-3, клип 1.0, 3000 шагов, batch 2.
Приёмка окна: held-ошибка на отложенных сэмплах (10 %), затем eval — генерация сетки и LPIPS/CLIP против базы.
Слияние в веса: W' = W + B·A·scale, поправка исчезает как отдельный объект, инференс не замедляется.

Окна. Семь слияний по два блока:

4,5 → 3     7,8 → 6     10,11 → 9     13,14 → 12     16,17 → 15     19,20 → 18     22,23 → 21

В артефакт пошли первые шесть (42 − 12 = 30 блоков). Седьмое окно даёт 28 блоков (combo7): CLIP 29.34 против 30.14 у combo6, и на глаз хуже — не используется. Восьмое (25,26 → 24) роняет CLIP до 25.4 — зона «громкая», слияние там не работает.

Почему поправки обучались в 1024 px. Первый прогон учил слияния на 512 px: CLIP на 512 был 30.14 против 29.49 у базы, но на 1024 падал до 28.92. Переобучение тех же окон на 1024 px дало 29.12 на 1024 (run_hi.sh — именно этот прогон); каждая сборка лучше «на своём» разрешении, в артефакт пошла версия для 1024.

Проверить, что всё работает

cd ..                                  # папка модели (там должен лежать исходный transformer/)
python3 src/window.py cache --start 3 --drop 2 --samples 2 --size 1024    # ~20 с
python3 src/window.py train --start 3 --drop 2 --samples 2 --size 1024 --steps 40   # ~10 с
python3 src/window.py combo --tags win3_5_1024,win6_8_1024 --tag check --size 1024 --sample-steps 10

combo соберёт модель из указанных слияний, сгенерит сетку по 10 промптам и посчитает LPIPS/CLIP против work/runs/base/ (её надо один раз создать: python3 src/lab.py baseline).

Что не сработало (чтобы не повторять)

Переподгонка на собранной модели (дообучить те же LoRA на собственных траекториях собранной модели, цели — исходная модель). Дрейф по границам блоков оказался 0.15–8 %, переподгонка его дожимала до 0.1–0.2 %, но CLIP падал 29.1 → 23.3 и лица разрушались. Метод лечил то, что почти не болело.
Текст в кадре слиянием не лечится: вывески и JSON-layout у 30-блочной модели деградируют (у базы «SALE 50% / only today», у сжатой — нечитаемая каша). На уровне состояний это не дрейф (2–8 %), значит потеря способности; нужен flow-matching на парах с надписями.
Сужение FFN (8192 → 6144, ещё −13 % параметров через закрытый пересчёт ffn_down) — не проверено, код в этой папке намеренно не оставлен.