#!/usr/bin/env python3
"""Файнтюн MOSS-Image DiT на закешированном датасете (`src/make_dataset.py`), цель — 16 ГБ.

Что грузится: только сам DiT и тензоры из датасета. Текстовый энкодер не нужен вообще —
условие и маска уже посчитаны, негатив CFG лежит рядом с датасетом в `empty_condition.npz`.
VAE нужен ровно для одного — посмотреть прогресс-сэмпл; по умолчанию он выключен
(`--sample-every 0`), и тогда процесс не трогает ничего, кроме датасета и DiT.

Почему влезает (замерено на этой машине, batch 1):
  * bf16-веса 6.94 ГиБ — база, от неё никуда не деться;
  * активации не держим: градиентный чекпойнтинг по блокам (`use_reentrant=False`), а входы
    блоков лежат в RAM и поднимаются на карту на пересчёт (~0.15 с/шаг шины);
  * градиенты живут на один блок: Adam применяется к блоку сразу после его backward
    (`BlockUpdater`), градиенты блока сразу освобождаются. Обычный Adam требует градиенты всех
    3.72 B одновременно (6.94 ГиБ) и на 1024x1024 падает с OOM при пике 14.7 ГиБ;
  * состояния Adam (2 байта × 2 на параметр = 7.45 ГиБ) живут в RAM — `PagedAdamW8bit`.
  Итог: пик 8.28 ГиБ на 768x1024 (2.2 с/шаг), 8.60 ГиБ на 1024x1024 (3.6 с/шаг).
  batch 2 нестабилен даже при 8 ГиБ запаса: managed-память paged-оптимизатора не успевает
  вытесняться, и torch получает OOM на 15.0 ГиБ — проверено, поэтому batch 1.

Про «fused forward/backward»: склеить forward и backward в один шаг нельзя — веса всё равно
нужны целиком. Память освобождает только «обновлять параметры сразу, как для них готов
градиент», и это здесь и сделано (BlockUpdater): пик падает с 14.7 до 8.6 ГиБ. Отличие от
обычного Adam одно — клип нормы считается по блоку, а не по всей модели.

Лосс — flow matching, как в планировщике модели:
    x_t = (1 - sigma) * x0 + sigma * noise,   таргет = noise - x0,   timestep = sigma * 1000.
Сетка обучения планировщика при `training_shift=1.0` равномерна по sigma, поэтому sigma — U(0,1).

Запуск (по умолчанию — одна эпоха по датасету):
    python3 src/train.py --steps 40                                  # смоук: 2 минуты
    python3 src/train.py --epochs 1 --save-every 2000                # прогон
    python3 src/train.py --epochs 2                                   # все датасеты в папке datasets

LR: `--lr` — базовая скорость; и прогрев, и затухание идут от `--lr/10`: линейный прогрев
lr/10 -> lr за `--warmup` шагов, затем cosine от lr до lr/10 к последнему шагу. Постоянного lr
нет — на сошедшейся базе он просто держит лосс.

Прогресс-сэмплы: раз в `--sample-every` секунд (по умолчанию час) из условий самого датасета
(сид фиксирован) -> `<out>/progress_<номер сэмпла>.png`; VAE на время сэмпла загружается и сразу
освобождается, на диске живут последние `--sample-keep` картинок.

Куда пишется результат: только в `--out` (по умолчанию `src/work/train`). Веса модели
(`transformer/`) скрипт только читает; если `--out` указывает внутрь компонентов модели
(`transformer/`, `vae/`, `text_encoder/`, `adapter/`, `tokenizer/`, `scheduler/`) или на саму папку
модели — падаем до загрузки весов, чтобы случайно не перезаписать модель. Один чекпойнт — ~7.4 ГБ.
"""
import argparse
import gc
import glob
import math
import os
import random as random_mod
import re
import shutil
import sys
import time

os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")

import numpy as np
import torch
import torch.nn.functional as F
from PIL import Image
from torch.utils.checkpoint import checkpoint

MODEL = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))  # папка модели
if MODEL not in sys.path:
    sys.path.insert(0, MODEL)  # diffusers ищет pipeline_moss_image как модуль

from pipeline_moss_image import MOSSImageScheduler, MOSSImageTransformer  # noqa: E402


def parse_args():
    p = argparse.ArgumentParser(description="MOSS-Image: файнтюн DiT под 16 ГБ")
    p.add_argument("--dataset", default=os.path.join(MODEL, "src", "work", "ds_testg_filtered"),
                   help="папка с датасетами и чанками из make_dataset.py; всё найденное внутри "
                        "склеивается в один набор")
    p.add_argument("--out", default=os.path.join(MODEL, "src", "work", "train_base"),
                   help="каталог логов, чекпойнтов и прогресс-сэмплов")
    p.add_argument("--epochs", type=float, default=1.0,
                   help="эпох по выбранным строкам; шаги = ceil(epochs * строк / (batch * accum))")
    p.add_argument("--steps", type=int, default=0,
                   help="точное число шагов, перебивает --epochs; 0 — считать из --epochs")
    p.add_argument("--lr", type=float, default=1e-5)
    p.add_argument("--warmup", type=int, default=200,
                   help="шагов линейного прогрева lr/10 -> lr")
    p.add_argument("--batch-size", type=int, default=1,
                   help="сэмплов на шаг (только одного размера); на 16 ГиБ влезает только 1 — "
                        "batch 2 доходит до 14.85 ГиБ и падает с OOM")
    p.add_argument("--accum", type=int, default=1, help="шагов накопления градиента")
    p.add_argument("--clip", type=float, default=1.0, help="клип нормы градиента; 0 — выключить")
    p.add_argument("--upd", choices=("block", "whole"), default="block",
                   help="block — Adam по блоку сразу после его backward (пик ~8 ГиБ); "
                        "whole — обычный Adam на всём наборе (нужно ~14.7 ГиБ, на 1024x1024 не влезает)")
    p.add_argument("--master-dtype", choices=("fp32", "bf16"), default="fp32",
                   help="в чём лежат мастер-веса для шага оптимизатора: fp32 — bf16-копия идёт в "
                        "forward/backward, а шаг делается по fp32 (иначе при lr~1e-5 обновление "
                        "теряется в bf16-решётке); bf16 — как было, шаг прямо по bf16-весам")
    p.add_argument("--master-device", choices=("cuda", "cpu"), default="cuda",
                   help="где живёт fp32-мастер: cuda — на карте (+13.9 ГиБ, шаг без копий через шину); "
                        "cpu — в RAM (на карте только bf16-веса, но градиенты блока едет по PCIe)")
    p.add_argument("--te", action="store_true", default=True,
                   help="считать условие на лету живым текстовым энкодером. Нужно, когда в датасете "
                        "нет колонок text_embeds/text_mask (make_dataset.py без --te); с колонками "
                        "энкодер не грузится")
    p.add_argument("--text-len", type=int, default=1024,
                   help="макс. токенов условия при подсчёте на лету (как в make_dataset.py)")
    p.add_argument("--word-dropout", type=float, default=0.05,
                   help="доля шагов, где из промпта выкидывается одно случайное слово (устойчивость "
                        "к ошибкам в описании); работает только когда энкодер загружен (--te или нет "
                        "колонок TE). 0 — выключить")
    p.add_argument("--text-dropout", type=float, default=0.05,
                   help="доля шагов с пустым промптом (CFG-ветка); без empty_condition.npz "
                        "молча выключается")
    p.add_argument("--sigma", choices=("uniform", "logit"), default="logit",
                   help="откуда берётся sigma: uniform — U(0,1), как учили базу; logit — "
                        "sigmoid(N(--sigma-shift, --sigma-scale)), смещает шаги к нижим сигмам (детали)")
    p.add_argument("--sigma-shift", type=float, default=-0.5,
                   help="logit-normal: среднее в logit-шкале; <0 — меньше шагов на высоком шуме, "
                        "больше на середине (детали)")
    p.add_argument("--sigma-scale", type=float, default=1.0, help="logit-normal: разброс")
    p.add_argument("--limit", type=int, default=0, help="первые N строк датасета (0 — все)")
    p.add_argument("--sizes", default="", help="только эти размеры, напр. 768x1024,1024x768")
    p.add_argument("--seed", type=int, default=42)
    p.add_argument("--log-every", type=int, default=50)
    p.add_argument("--ema", type=float, default=0.99,
                   help="сглаживание лосса в логе: ema = ema*--ema + loss*(1---ema); 0 — без сглаживания")
    p.add_argument("--save-every", type=int, default=2000,
                   help="шагов между чекпойнтами; 0 — только в конце")
    p.add_argument("--keep", type=int, default=4,
                   help="сколько step_* чекпойнтов держать на диске (старые удаляются)")
    p.add_argument("--init", default="", help="стартовать с чекпойнта (step_0005000 или final)")
    p.add_argument("--start-step", type=int, default=0,
                   help="с какого шага продолжать расписание lr при --init (иначе прогрев и cosine "
                        "пойдут заново с нуля — неверно при продолжении)")
    p.add_argument("--sample-every", type=float, default=1800.0,
                   help="секунд между прогресс-сэмплами (1800 — раз в полчаса); 0 — выключить "
                        "(тогда VAE не грузится)")
    p.add_argument("--sample-keep", type=int, default=8,
                   help="сколько progress_*.png держать на диске (старые удаляются); 0 — не удалять")
    p.add_argument("--sample-n", type=int, default=6, help="картинок в прогресс-сэмпле")
    p.add_argument("--sample-steps", type=int, default=30, help="шагов денойзинга в сэмпле")
    p.add_argument("--sample-size", type=int, default=1024, help="длинная сторона сэмпла")
    p.add_argument("--guidance", type=float, default=3.0)
    p.add_argument("--sample-rows", default="",
                   help="индексы строк для прогресс-сэмплов через запятую, напр. 0,30000; "
                        "по умолчанию — равномерно по датасету")
    p.add_argument("--sample-seed", type=int, default=42)
    return p.parse_args()


# ---------------------------------------------------------------- данные
def find_dataset_chunks(root):
    """Все arrow-датасеты под `root` (папки с dataset_info.json).

    make_dataset.py пишет не один датасет, а кучу под-датасетов `chunk_*/sub_*` — их и находим.
    """
    chunks = []
    for dirpath, dirnames, filenames in os.walk(root):
        dirnames[:] = [d for d in dirnames if not d.startswith(".")]
        if "dataset_info.json" in filenames or "state.json" in filenames:
            chunks.append(dirpath)
            dirnames.clear()  # внутрь датасета не лезем
    return chunks


def load_any_dataset(path):
    """Один датасет или папка с чанками: конкатенация ленивая, arrow-файлы не переписываются."""
    from datasets import concatenate_datasets, load_from_disk
    chunks = find_dataset_chunks(path)
    if not chunks:
        raise FileNotFoundError(f"в {path} нет arrow-датасета (ни dataset_info.json, ни state.json)")
    if len(chunks) == 1:
        print(f"датасет: {chunks[0]}", flush=True)
        return load_from_disk(chunks[0])
    print(f"датасет: {len(chunks)} чанков в {path}", flush=True)
    return concatenate_datasets([load_from_disk(c) for c in chunks])


def open_dataset(path, limit, sizes):
    """Датасет не грузим в RAM: латент на строку — мегабайты, а 6k строк — это десятки ГБ.

    Возвращаем (ds, rows): `rows` — только метаданные (индекс, размер), сами тензоры
    читаются по индексу в `to_batch` (27 мс на строку при 768x1024, незаметно на фоне шага).
    """
    ds = load_any_dataset(path)
    meta = ds.select_columns(["width", "height"])  # только скаляры, быстро
    ds = ds.with_format("numpy")
    wanted = {x.strip() for x in sizes.split(",") if x.strip()}
    rows = []
    for i in range(len(meta)):
        w, h = int(meta[i]["width"]), int(meta[i]["height"])
        if wanted and f"{w}x{h}" not in wanted:
            continue
        rows.append({"idx": i, "width": w, "height": h})
        if limit > 0 and len(rows) >= limit:  # limit считается после фильтра по размерам
            break
    return ds, rows


def load_empty_condition(path):
    """Негатив CFG, посчитанный make_dataset.py, или None.

    Ищем и в корне, и на уровень ниже: при `--dataset datasets` файл лежит в каждом под-датасете
    (`datasets/civit/empty_condition.npz`), а в корне его нет — без этого сэмплы шли без CFG.
    """
    files = [os.path.join(path, "empty_condition.npz")]
    files += sorted(glob.glob(os.path.join(path, "*", "empty_condition.npz")))
    for file in files:
        if os.path.exists(file):
            z = np.load(file)
            return (torch.from_numpy(z["embeds"].copy()), torch.from_numpy(z["mask"].copy()))
    return None


class Buckets:
    """Батчи только из одинаковых размеров: латент — один тензор, разные HxW не склеить."""

    def __init__(self, rows, generator):
        self.rows = rows
        self.generator = generator
        self.by_size = {}
        for i, r in enumerate(rows):
            self.by_size.setdefault((r["width"], r["height"]), []).append(i)
        self.sizes = sorted(self.by_size)
        print("размеры: " + ", ".join(f"{w}x{h}: {len(self.by_size[(w, h)])}"
                                      for w, h in self.sizes), flush=True)
        self._perms, self.cursor = {}, {}

    def _shuffle(self):
        for size, idx in self.by_size.items():
            order = list(idx)
            for i in range(len(order) - 1, 0, -1):  # фишер-йетс на своём генераторе
                j = int(torch.randint(0, i + 1, (1,), generator=self.generator))
                order[i], order[j] = order[j], order[i]
            self._perms[size] = order
            self.cursor[size] = 0

    def batch(self, batch_size):
        if not self._perms:
            self._shuffle()
        # размер выбираем пропорционально числу строк, чтобы мелкие группы не залипали
        weights = [len(self.by_size[s]) for s in self.sizes]
        pick = int(torch.multinomial(torch.tensor(weights, dtype=torch.float), 1,
                                     generator=self.generator))
        size = self.sizes[pick]
        picked = []
        for _ in range(batch_size):
            if self.cursor[size] >= len(self._perms[size]):
                self._perms[size] = self._perms[size][::-1]  # второй проход в обратном порядке
                self.cursor[size] = 0
            picked.append(self._perms[size][self.cursor[size]])
            self.cursor[size] += 1
        return [self.rows[i] for i in picked]


def read_row(ds, row):
    r = ds[row["idx"]]
    return {"vae": r["vae"], "text": r["text"], "text_embeds": r.get("text_embeds"),
            "text_mask": r.get("text_mask"), "source": r["source"], "width": row["width"],
            "height": row["height"]}


def to_batch(rows, ds, device):
    """rows (одинаковый HxW) -> латент [B,128,h,w] bf16, условие [B,L,4096] bf16 + маска [B,L].

    Если в датасете нет text_embeds (make_dataset.py без --te), условие собирается живым
    энкодером в text_encoder.py — возвращаем строки, а не тензоры.
    """
    recs = [read_row(ds, r) for r in rows]
    lat = torch.stack([torch.from_numpy(r["vae"]) for r in recs]).to(device, torch.bfloat16)
    if recs[0]["text_embeds"] is None:
        return lat, None, None
    length = max(len(r["text_embeds"]) for r in recs)
    cond = torch.zeros(len(recs), length, 4096, dtype=torch.bfloat16)
    mask = torch.zeros(len(recs), length, dtype=torch.bool)
    for i, r in enumerate(recs):
        n = len(r["text_embeds"])
        cond[i, :n] = torch.from_numpy(r["text_embeds"])
        mask[i, :n] = torch.from_numpy(r["text_mask"])
    return lat, cond.to(device), mask.to(device)


# ---------------------------------------------------------------- условие на лету
class LiveTextEncoder:
    """Текст условия живым ТЕ, когда в датасете нет колонок text_embeds/text_mask.

    Путь тот же, что в make_dataset.py (`encode_text`/`render`), чтобы условие совпало с тем,
    что было бы закэшировано: chat-шаблон с add_generation_prompt=True и enable_thinking=True,
    обрезка по --text-len, отдельный forward на промпт без паддинга (паддинг менял бы условие).
    Возвращает (embeds float16 [L,4096], mask bool [L]).
    """

    def __init__(self, model_dir, device, text_len):
        """`device` — где считать; на карте энкодер держится только на время кодирования.

        0.85B стоит ~2.3 ГиБ. Вместе с 20.8 ГиБ fp32-мастера и активациями DiT на 32 ГБ это
        упирается в потолок, поэтому после encode_batch энкодер уходит в offload (см. `release`).
        """
        from transformers import AutoTokenizer
        from pipeline_moss_image import MOSSAdapterTextEncoder
        path = os.path.join(model_dir, "text_encoder")
        self.tok = AutoTokenizer.from_pretrained(path)
        # from_pretrained сам поднимает голову из ../adapter
        self.te = MOSSAdapterTextEncoder.from_pretrained(path).to(device).eval()
        self.te.requires_grad_(False)
        self.device = torch.device(device)
        self.offloaded = False
        self.text_len = text_len

    def release(self):
        """Убрать энкодер с карты, если он там: память нужна шагу DiT."""
        if self.device.type == "cuda" and not self.offloaded:
            self.te.to("cpu")
            self.offloaded = True
            torch.cuda.empty_cache()

    def _ensure(self):
        if self.offloaded:
            self.te.to(self.device)
            self.offloaded = False

    def _render(self, text):
        return self.tok.apply_chat_template([{"role": "user", "content": text}], tokenize=False,
                                            add_generation_prompt=True, enable_thinking=True)

    @torch.no_grad()
    def encode(self, text):
        self._ensure()
        ids = self.tok(self._render(text), truncation=True, max_length=self.text_len)["input_ids"]
        if not ids:
            raise ValueError("пустой промпт после токенизации")
        ids = torch.tensor([ids], dtype=torch.long, device=self.device)
        out = self.te(input_ids=ids, attention_mask=torch.ones_like(ids, dtype=torch.bool),
                      use_cache=False)
        emb = out.last_hidden_state[0].to(torch.float16)
        return emb.cpu().numpy(), np.ones(emb.shape[0], dtype=bool)

    def encode_batch(self, texts):
        """Строки промптов -> [B,L,4096] bf16 + маска [B,L] (L = max по батчу, без паддинга в ТЕ)."""
        pairs = [self.encode(t) for t in texts]
        length = max(len(e) for e, _ in pairs)
        cond = torch.zeros(len(pairs), length, 4096, dtype=torch.bfloat16)
        mask = torch.zeros(len(pairs), length, dtype=torch.bool)
        for i, (emb, m) in enumerate(pairs):
            cond[i, :len(emb)] = torch.from_numpy(emb)
            mask[i, :len(emb)] = torch.from_numpy(m)
        return cond, mask


def drop_word(text, gen, rng):
    """Убрать одно случайное слово из промпта (один бросок на строку).

    Цель — чтобы модель была устойчива к пропускам и мусору в описании. Если слово одно —
    остаётся пустая строка (это же и есть ветка пустого промпта).
    """
    words = text.split()
    if not words:
        return text
    del words[int(torch.randint(0, len(words), (1,), generator=gen))]
    out = " ".join(words)
    return out or rng.choice(["image", "a picture", "photo"])


def _pair(module, hidden_dtype, master_dtype, device):
    """Завести мастер-тензоры для Linear: (master_w, master_b), скрытые веса — в module.weight."""
    if not isinstance(module, torch.nn.Linear):
        return
    def empty_like(t):
        return torch.empty(t.shape, dtype=master_dtype, device=device)
    pairs = [empty_like(module.weight)]
    if module.bias is not None:
        pairs.append(empty_like(module.bias))
    module._master = pairs


class MasterWeights:
    """fp32-мастер весов: forward/backward по bf16, шаг оптимизатора — по fp32-копии.

    Зачем: при lr ~1e-5 обновление `w -= lr*g` — это изменение ~1e-5 от |w|, а квант шага bf16
    (7 мантийных бит) ~2e-3 от |w|. Шаг не представим: он либо равен нулю, либо округляется
    вверх до ступеньки решётки. Мастер в fp32 убирает именно это; веса для forward остаются bf16,
    поэтому VRAM и скорость счёта не меняются.

    Устройство: `module.weight` всегда остаётся bf16 (`nn.Parameter` не подменяется — иначе
    autograd и оптимизатор начинают спорить о dtype градиента). Рядом лежит fp32-тензор
    `module._master[i]`, уже объявленный оптимизатору как параметр. Цикл шага:

        forward/backward   по bf16-весам
        grads()            bf16-градиенты -> fp32-grad мастеров
        opt.step()         Adam в fp32 по мастерам (состояния тоже fp32)
        sync()             мастера -> bf16-веса для следующего forward

    Предел: fp32-мастер на всю модель — это +13.9 ГиБ (на карте или в RAM при --master-device cpu).
    """

    def __init__(self, model, device, dtype=torch.float32):
        self.device = torch.device(device)
        self.master_dtype = dtype
        self.slots = []  # [(bf16-параметр, fp32-мастер)]
        self._map = None
        hidden_dtype = next(model.parameters()).dtype
        for module in model.modules():
            _pair(module, hidden_dtype, dtype, self.device)
        for module in model.modules():
            masters = getattr(module, "_master", None)
            if masters is None:
                continue
            hidden = [module.weight] + ([module.bias] if module.bias is not None else [])
            self.slots += list(zip(hidden, masters))

    @torch.no_grad()
    def init(self, model):
        """Один раз, до создания оптимизатора: мастера = стартовые веса в fp32."""
        for hidden, master in self.slots:
            master.copy_(hidden.data.to(self.master_dtype))

    def parameters(self):
        return [master for _, master in self.slots]

    def _master_of(self, hidden):
        """fp32-двойник bf16-параметра (словарь строится один раз)."""
        if self._map is None:
            self._map = {id(h): m for h, m in self.slots}
        return self._map.get(id(hidden))

    @torch.no_grad()
    def grads(self):
        """bf16-градиенты -> fp32-градиенты мастеров (None, если градиента нет)."""
        self.grads_for([h for h, _ in self.slots])

    @torch.no_grad()
    def grads_for(self, hidden_params):
        """То же для конкретной группы (нужно BlockUpdater, который обновляет по блоку)."""
        for h in hidden_params:
            m = self._master_of(h)
            if m is not None:
                m.grad = None if h.grad is None else h.grad.to(m.device, self.master_dtype)

    @torch.no_grad()
    def sync(self):
        """fp32-мастер -> bf16-веса для следующего forward."""
        for hidden, master in self.slots:
            hidden.data.copy_(master)

    @torch.no_grad()
    def sync_one(self, hidden):
        """fp32-мастер одного параметра -> его bf16-вес."""
        m = self._master_of(hidden)
        if m is not None:
            hidden.data.copy_(m)

# ---------------------------------------------------------------- чекпойнтинг и сэмплы
def _run_block(block, x_cpu, rope, mod):
    """Вход блока лежит на CPU, на пересчёт поднимается обратно (см. set_block_checkpointing)."""
    return block(x_cpu.to("cuda", non_blocking=True), rope, mod)


def set_block_checkpointing(tr, on):
    """Пересчёт блоков на backward вместо хранения активаций — без этого 16 ГБ не хватает.

    Вход каждого блока (25 МиБ на 1024 px) держим в RAM, а не на карте: autograd сохраняет
    только этот вход, поэтому выигрываем ~0.75-1.0 ГиБ GPU — ровно ту разницу, которой не
    хватало на квадратных 1024x1024 и на управляемую память PagedAdamW. За шаг это ~0.15 с
    копий шины, дешевле любого другого способа освободить гигабайт.
    """
    for layer in tr.layers:
        if on and "forward" not in layer.__dict__:
            block = layer.forward
            layer.forward = (lambda x, rope, mod, _b=block:
                             checkpoint(_run_block, _b, x.cpu(), rope, mod, use_reentrant=False))
        elif not on and "forward" in layer.__dict__:
            del layer.forward


class BlockUpdater:
    """Adam по каждому блоку сразу после его backward; градиенты блока освобождаются.

    Зачем: обычному Adam нужны градиенты всех 3.72 B параметров (6.94 ГиБ) одновременно —
    вместе с весами это 13.88 ГиБ, и 1024x1024 в 16 ГБ уже не влезает (замерено: OOM на 14.7).
    `register_post_accumulate_grad_hook` срабатывает, когда у параметра появился градиент; как
    только все параметры блока их получили, блок обновляется и его градиенты обнуляются.
    Пик: веса + активации одного блока ~ 8 ГиБ.

    Градиенты при этом те же, что в обычном проходе: backward идёт от последнего блока к первому,
    и блок обновляется только после того, как посчитан и пересчитан на своих старых весах.
    Отличие одно — клип локальный по блоку, а не глобальный (это единственное, что нельзя
    посчитать без всех градиентов сразу).
    """

    def __init__(self, tr, lr, clip, master=None):
        import bitsandbytes as bnb
        block_ids = {id(p) for layer in tr.layers for p in layer.parameters()}
        groups = [[p for p in layer.parameters()] for layer in tr.layers]
        groups.append([p for p in tr.parameters() if id(p) not in block_ids])  # эмбеддеры, голова
        self.groups = [g for g in groups if g]
        self.clip = clip
        # с fp32-мастером шаг идёт по мастерам: у каждого bf16-параметра свой fp32-двойник
        self.master = master
        to_master = {} if master is None else {id(h): m for h, m in master.slots}
        self.targets = [[to_master.get(id(p), p) for p in g] for g in self.groups]
        self.opts = [bnb.optim.PagedAdamW8bit(t, lr=lr, betas=(0.9, 0.95), eps=1e-8)
                     for t in self.targets]
        self.left = []
        for gi, params in enumerate(self.groups):
            self.left.append([len(params)])
            for p in params:
                p.register_post_accumulate_grad_hook(lambda _p, gi=gi: self._release(gi))

    def set_lr(self, lr):
        for opt in self.opts:
            for group in opt.param_groups:
                group["lr"] = lr

    def _release(self, gi):
        self.left[gi][0] -= 1
        if self.left[gi][0] > 0:
            return
        params = self.targets[gi]
        if self.master is not None:
            self.master.grads_for(self.groups[gi])
        if self.clip > 0:
            torch.nn.utils.clip_grad_norm_(params, self.clip)
        self.opts[gi].step()
        for p in params:
            p.grad = None
        for p in self.groups[gi]:  # bf16-копии получают обновлённые веса
            if self.master is not None:
                self.master.sync_one(p)
            p.grad = None
        self.left[gi][0] = len(params)

    def finish(self):
        """Страховка: группа, которой не хватило градиентов, обновляется здесь."""
        for gi, left in enumerate(self.left):
            params = self.groups[gi]
            if left[0] == len(params):
                continue
            print(f"ВНИМАНИЕ: группе {gi} не хватило {left[0]} градиентов — обновляю в конце",
                  flush=True)
            if any(p.grad is not None for p in params):
                if self.clip > 0:
                    torch.nn.utils.clip_grad_norm_(params, self.clip)
                self.opts[gi].step()
            for p in params:
                p.grad = None
            left[0] = len(params)


def prune_samples(out, keep):
    """Имена нумерованы по сэмплу и с ведущими нулями, поэтому старые — это просто начало списка."""
    if keep <= 0:
        return
    names = sorted(n for n in os.listdir(out)
                   if n.startswith("progress_") and n.endswith(".png"))
    for name in names[:-keep]:
        os.remove(os.path.join(out, name))


def sample_grid(tr, rows, ds, args, out_path, empty, live=None):
    """Несколько картинок из условий датасета: VAE грузим на время сэмпла и сразу освобождаем."""
    from diffusers import AutoencoderKLFlux2
    from pipeline_moss_image import MOSSImagePipeline
    set_block_checkpointing(tr, False)
    tr.eval()
    # fp16: замерено, что на этом VAE качество то же, что в fp32, а грузится и считает вдвое быстрее
    vae = AutoencoderKLFlux2.from_pretrained(os.path.join(MODEL, "vae"),
                                             torch_dtype=torch.float16).to("cuda").eval()
    pipe = MOSSImagePipeline(tr, MOSSImageScheduler(), None, None, vae)
    pipe.set_progress_bar_config(disable=True)
    cfg = args.guidance > 1.0 and (empty is not None or live is not None)
    if cfg:
        if live is not None:
            e, m = live.encode("")
            nemb = torch.from_numpy(e)[None].to("cuda", torch.bfloat16)
            nmask = torch.from_numpy(m)[None].cuda()
        else:
            nemb, nmask = empty[0][None].cuda(), empty[1][None].cuda()
    tiles = []
    for i, meta in enumerate(rows):
        r = read_row(ds, meta)
        long_side = args.sample_size  # aspect из датасета, длинная сторона = --sample-size
        if r["width"] >= r["height"]:
            w, h = long_side, round(long_side * r["height"] / r["width"] / 32) * 32
        else:
            w, h = round(long_side * r["width"] / r["height"] / 32) * 32, long_side
        print(f"  сэмпл {i + 1}: {w}x{h}  {r['text'][:200]}", flush=True)
        if r["text_embeds"] is not None:
            cond = torch.from_numpy(r["text_embeds"])[None].to("cuda", torch.bfloat16)
            cmask = torch.from_numpy(r["text_mask"])[None].cuda()
        else:  # условия в датасете нет — считаем тем же путём, что и обучение
            cond, cmask = live.encode_batch([r["text"]])
            cond, cmask = cond.to("cuda"), cmask.to("cuda")
        tiles.append(pipe(
            prompt_embeds=cond,
            prompt_attention_mask=cmask,
            negative_prompt_embeds=nemb if cfg else None,
            negative_prompt_attention_mask=nmask if cfg else None,
            height=h, width=w, num_inference_steps=args.sample_steps,
            guidance_scale=args.guidance if cfg else 1.0,
            generator=torch.Generator("cuda").manual_seed(args.sample_seed + i),
            output_type="pil").images[0])
    del pipe, vae
    gc.collect()
    torch.cuda.empty_cache()

    cols = min(2, len(tiles))
    rows_n = (len(tiles) + cols - 1) // cols
    cell = args.sample_size
    grid = Image.new("RGB", (cols * cell, rows_n * cell), "white")
    for i, tile in enumerate(tiles):  # ячейка фиксированная, картинка по центру — без искажений
        grid.paste(tile, ((i % cols) * cell + (cell - tile.width) // 2,
                          (i // cols) * cell + (cell - tile.height) // 2))
    grid.save(out_path)
    print(f"сэмпл -> {out_path}" + ("" if cfg else " (без CFG: нет empty_condition.npz)"),
          flush=True)
    tr.train()
    set_block_checkpointing(tr, True)  # без этого forward DiT держит все активации и ловит OOM


# ---------------------------------------------------------------- обучение
PROTECTED = ("transformer", "vae", "text_encoder", "adapter", "tokenizer", "scheduler")


def guard_out(out):
    """Не даём писать чекпойнты поверх модели: обучение — только в отдельный каталог."""
    out_abs = os.path.abspath(out)
    model_abs = os.path.abspath(MODEL)
    for name in [""] + list(PROTECTED):
        protected = os.path.join(model_abs, name) if name else model_abs
        if out_abs == protected or out_abs.startswith(protected + os.sep) and name:
            raise SystemExit(f"--out {out_abs} лежит внутри модели ({protected}); "
                             f"выберите отдельный каталог, например src/work/train")


def main():
    args = parse_args()
    guard_out(args.out)
    torch.manual_seed(args.seed)
    os.makedirs(args.out, exist_ok=True)
    device = "cuda"

    # PagedAdamW8bit не имеет своего лимита: вытесненные состояния возвращаются в VRAM, как только
    # там освободилось место, карта доходит до нуля запаса и падает OOM на транзиентной аллокации
    # (40 МиБ на пересчёте FFN). Потолок аллокатора оставляет ~8 % свободными; не влезшие страницы
    # оптимизатор держит в RAM, это его штатный режим.
    torch.cuda.set_per_process_memory_fraction(0.92)
    print(f"потолок VRAM: {0.92 * torch.cuda.get_device_properties(0).total_memory / 2**30:.2f} ГиБ "
          f"из {torch.cuda.get_device_properties(0).total_memory / 2**30:.2f}", flush=True)

    ds, rows = open_dataset(args.dataset, args.limit, args.sizes)
    if not rows:
        raise SystemExit("датасет пуст")
    print(f"строк: {len(rows)}", flush=True)

    has_te_cols = "text_embeds" in ds.column_names
    live = None
    if args.te or not has_te_cols:
        live = LiveTextEncoder(MODEL, device, args.text_len)
        print(f"условие считается на лету ({'--te' if args.te else 'в датасете нет text_embeds'}), "
              f"text-len {args.text_len}", flush=True)
    empty = load_empty_condition(args.dataset)
    if live is None and empty is not None:
        print(f"пустой промпт (негатив CFG): {tuple(empty[0].shape)}", flush=True)
    if live is None and args.text_dropout > 0 and empty is None:
        print("empty_condition.npz нет -> text-dropout выключен", flush=True)
    dropout = args.text_dropout if (live is not None or empty is not None) else 0.0
    # дропаут слова требует перекодирования промпта: работает только когда энкодер загружен
    word_dropout = args.word_dropout if live is not None else 0.0
    if args.word_dropout > 0 and live is None:
        print("word-dropout выключен: условие берётся из датасета, энкодер не загружен", flush=True)
    if word_dropout:
        print(f"дропауты: пустой промпт {dropout:.0%}, случайное слово {word_dropout:.0%}",
              flush=True)

    n_sample = min(args.sample_n, len(rows))  # сэмплы — из самих данных, по всему датасету
    if args.sample_rows:  # при смешанных датасетах равномерный срез может попасть в tag-промпт
        sample_rows = [rows[int(x)] for x in args.sample_rows.split(",") if x.strip()]
    else:
        sample_rows = [rows[int(i)] for i in np.linspace(0, len(rows) - 1, n_sample)]
    if args.sample_every:  # промпты в лог сразу: иначе непонятно, что именно нарисовано
        text_col = ds.select_columns(["text"])
        print("прогресс-сэмплы берутся из строк:", flush=True)
        for i, meta in enumerate(sample_rows, 1):
            print(f"  {i}) строка {meta['idx']} {meta['width']}x{meta['height']}  "
                  f"{text_col[meta['idx']]['text'][:160]}", flush=True)

    # Paged-оптимизатор держит состояния в managed-памяти: если карту уже занял другой процесс,
    # он не падает с OOM, а часами гоняет страницы. Лучше предупредить сразу.
    free_gib = torch.cuda.mem_get_info()[0] / 2**30
    if free_gib < 15.0:
        print(f"ВНИМАНИЕ: свободно {free_gib:.1f} ГиБ — обучению нужно ~14.7. "
              f"Проверьте, не висит ли другой процесс на GPU (nvidia-smi)", flush=True)
    tr = MOSSImageTransformer.from_pretrained(
        os.path.join(MODEL, "transformer"), torch_dtype=torch.bfloat16,
        local_files_only=True, low_cpu_mem_usage=False).to(device).train()
    if args.init:
        state = MOSSImageTransformer.from_pretrained(
            args.init, torch_dtype=torch.bfloat16, local_files_only=True,
            low_cpu_mem_usage=False).state_dict()
        tr.load_training_state_dict(state)
        print(f"старт с {args.init}", flush=True)
    tr.requires_grad_(True)
    params = [p for p in tr.parameters() if p.requires_grad]
    n_param = sum(p.numel() for p in params)
    print(f"параметров {n_param / 1e9:.3f}B, веса bf16 {n_param * 2 / 2**30:.2f} ГиБ, "
          f"градиенты — только текущего блока", flush=True)
    master = None
    if args.master_dtype == "fp32":
        # до создания оптимизатора: Adam должен видеть fp32-параметры, а не bf16
        master = MasterWeights(tr, args.master_device)
        master.init(tr)
        extra = f"+{n_param * 4 / 2**30:.2f} ГиБ на {args.master_device}"
        print(f"мастер-веса fp32 ({extra}), forward/backward остаются bf16", flush=True)
    else:
        print("мастер-веса bf16: шаг оптимизатора идёт прямо по bf16-весам", flush=True)
    set_block_checkpointing(tr, True)

    if args.upd == "block":
        if args.accum != 1:
            raise SystemExit("--upd block требует --accum 1: обновление идёт внутри backward "
                             "и не может ждать следующего микро-батча")
        updater, opt = BlockUpdater(tr, args.lr, args.clip, master), None
        print(f"оптимизатор: block-wise PagedAdamW8bit ({len(updater.groups)} групп, "
              f"градиенты живут только на один блок)", flush=True)
    else:
        import bitsandbytes as bnb
        updater = None
        targets = master.parameters() if master is not None else params
        opt = bnb.optim.PagedAdamW8bit(targets, lr=args.lr, betas=(0.9, 0.95), eps=1e-8)
        print("оптимизатор: общий PagedAdamW8bit "
              f"({'fp32-мастера' if master is not None else 'bf16-веса'})", flush=True)

    gen = torch.Generator().manual_seed(args.seed)
    buckets = Buckets(rows, gen)
    if args.sigma == "logit":  # печатаем, что именно получилось на бумаге
        probe = torch.sigmoid(torch.randn(200000) * args.sigma_scale + args.sigma_shift)
        print(f"sigma: logit-normal m={args.sigma_shift:+.2f} s={args.sigma_scale:.2f}, "
              f"медиана {probe.median():.2f}, доля шагов sigma<0.2: {(probe < 0.2).float().mean():.0%}, "
              f"sigma>0.8: {(probe > 0.8).float().mean():.0%}", flush=True)
    else:
        print("sigma: uniform U(0,1), как в обучении базы", flush=True)
    # «эпоха» здесь — один проход по выбранным строкам: перемешивание идёт внутри размера,
    # поэтому порядок не строгий, но среднее число просмотров на строку верное
    step_rows = args.batch_size * args.accum
    new_steps = args.steps or max(1, math.ceil(args.epochs * len(rows) / step_rows))
    start_step = args.start_step  # уже пройденные шаги: lr продолжается, а не стартует заново
    steps = new_steps + start_step
    min_lr = args.lr / 10
    print(f"шагов: {new_steps}" + (f" ({args.epochs:g} эпох x {len(rows)} строк / {step_rows})"
                             if not args.steps else " (--steps)")
          + (f", продолжаю с {start_step} -> до {steps}" if start_step else "")
          + f", lr {args.lr:.1e} -> {min_lr:.1e}, warmup {args.warmup}", flush=True)
    log_path = os.path.join(args.out, "log.csv")
    with open(log_path, "a" if start_step else "w") as log:
        if not start_step:
            log.write("step,loss_ema,loss_raw,lr,sec,peak_gib,eta_h\n")

    def lr_at(step):
        """Прогрев lr/10 -> lr, затем cosine lr -> lr/10. Постоянного lr нет."""
        if args.warmup > 0 and step < args.warmup:
            return min_lr + (args.lr - min_lr) * (step + 1) / args.warmup
        progress = min(1.0, max(0.0, (step - args.warmup) / max(1, steps - args.warmup)))
        return min_lr + (args.lr - min_lr) * 0.5 * (1 + math.cos(math.pi * progress))

    ema = None  # сглаженный лосс: по сырым значениям ничего не видно (sigma у каждого шага своя)
    sec_sum = 0.0  # сумма времени шагов — для ETA по фактическому среднему
    t_start = time.time()
    t_last_sample = time.time()
    sample_n = 0  # в имя файла идёт номер сэмпла, а не шаг: тогда старые видно и можно удалять
    torch.cuda.reset_peak_memory_stats()
    if args.sample_every:  # «до обучения» — чтобы было с чем сравнивать
        sample_grid(tr, sample_rows, ds, args,
                    os.path.join(args.out, "progress_00000.png"), empty, live)
        sample_n += 1
    for step in range(start_step + 1, steps + 1):
        lr = lr_at(step - 1)
        if updater is not None:
            updater.set_lr(lr)
        else:
            for group in opt.param_groups:
                group["lr"] = lr
        t0 = time.time()
        total = 0.0
        for _ in range(args.accum):
            batch = buckets.batch(args.batch_size)
            lat, cond, mask = to_batch(batch, ds, device)
            if live is not None:
                texts = [read_row(ds, r)["text"] for r in batch]
                dropped = False
                # два независимых броска: 1) пустой промпт (CFG-ветка), 2) одно случайное слово
                if dropout and torch.rand(1, generator=gen).item() < dropout:
                    texts, dropped = [""] * len(batch), True
                elif word_dropout and torch.rand(1, generator=gen).item() < word_dropout:
                    i = int(torch.randint(0, len(batch), (1,), generator=gen))
                    texts[i] = drop_word(texts[i], gen, random_mod)
                    dropped = True
                if dropped or cond is None:  # без дропаута условие из датасета не трогаем
                    cond, mask = live.encode_batch(texts)
                    cond, mask = cond.to(device), mask.to(device)
                live.release()  # энкодер с карты: память нужна шагу DiT
                del texts, dropped
            noise = torch.randn(lat.shape, dtype=torch.float32, device=device)
            if args.sigma == "logit":  # больше шагов там, где модель доводит детали (низкие sigma)
                sigma = torch.sigmoid(torch.randn(len(batch), 1, 1, 1, device=device)
                                      * args.sigma_scale + args.sigma_shift)
            else:
                sigma = torch.rand(len(batch), 1, 1, 1, device=device)
            x0 = lat.float()
            xt = ((1 - sigma) * x0 + sigma * noise).to(torch.bfloat16)
            target = noise - x0
            pred = tr(xt, sigma.flatten() * 1000.0, cond, mask).sample.float()
            loss = F.mse_loss(pred, target) / args.accum
            loss.backward()
            total += loss.item()
            del lat, cond, mask, xt, target, pred, loss
        if updater is not None:
            updater.finish()  # в норме это no-op: все группы уже обновлены хуками
        else:
            if master is not None:
                master.grads()  # bf16-градиенты -> fp32-градиенты мастеров
                params_to_clip = master.parameters()
            else:
                params_to_clip = params  # параметры в --init не попали в master.slots
            if args.clip > 0:
                torch.nn.utils.clip_grad_norm_(params_to_clip, args.clip)
            opt.step()
            opt.zero_grad(set_to_none=True)
            if master is not None:
                master.sync()  # fp32-мастер -> bf16-веса для следующего forward
        step_sec = time.time() - t0
        sec_sum += step_sec  # средняя по шагам, без сохранений чекпойнтов и сэмплов
        peak = torch.cuda.max_memory_allocated() / 2**30
        raw = total / args.accum
        ema = raw if ema is None or args.ema <= 0 else args.ema * ema + (1 - args.ema) * raw

        if step % args.log_every == 0 or step == 1:
            done = step - start_step
            eta = sec_sum / max(done, 1) * (steps - step)  # среднее по шагам ЭТОГО прогона
            eta_h = eta / 3600
            print(f"шаг {step}/{steps}  loss {ema:.4f} (raw {raw:.3f})  lr {lr:.1e}  "
                  f"{step_sec:.1f} c/шаг  пик {peak:.2f} ГиБ  "
                  f"осталось {int(eta_h)}ч{int(eta % 3600 // 60):02d}м", flush=True)
            with open(log_path, "a") as log:
                log.write(f"{step},{ema:.6f},{raw:.6f},{lr:.2e},{step_sec:.2f},{peak:.2f},"
                          f"{eta_h:.2f}\n")

        if args.save_every and step % args.save_every == 0:
            path = os.path.join(args.out, f"step_{step:05d}")
            tr.save_pretrained(path)
            print(f"чекпойнт -> {path}", flush=True)
            old = sorted((d for d in os.listdir(args.out) if d.startswith("step_")),
                         key=lambda d: os.path.getmtime(os.path.join(args.out, d)))
            # по времени, а не по номеру: после перезапуска с --init в папке остаются чекпойнты
            # прошлого прогона с большими номерами, и сортировка по номеру удаляла свежие
            for name in old[:-args.keep] if args.keep > 0 else []:
                shutil.rmtree(os.path.join(args.out, name), ignore_errors=True)
                print(f"старый чекпойнт удалён: {name}", flush=True)

        if args.sample_every and (time.time() - t_last_sample >= args.sample_every
                                  or step == steps):
            sample_grid(tr, sample_rows, ds, args,
                        os.path.join(args.out, f"progress_{sample_n:05d}.png"), empty, live)
            sample_n += 1
            t_last_sample = time.time()
            prune_samples(args.out, args.sample_keep)

    tr.save_pretrained(os.path.join(args.out, "final"))
    peak = torch.cuda.max_memory_allocated() / 2**30
    print(f"готово: {steps} шагов за {time.time() - t_start:.0f} c, пик VRAM {peak:.2f} ГиБ, "
          f"чекпойнт {os.path.join(args.out, 'final')}, лог {log_path}", flush=True)
    if peak > 15.5:
        print("ВНИМАНИЕ: пик близко к 16 ГиБ — на другой карте шаг может не влезть", flush=True)


if __name__ == "__main__":
    main()