Blame for src/train.py

#!/usr/bin/env python3
"""Обучение UDiT на папке «картинка + .txt». Квадратные кадры (--size), кэш латентов и текста.

    python3 src/train.py --data ~/dataset/butterflies_256sq/train \
            --eval ~/dataset/butterflies_256sq/eval --epochs 80 --batch 8

Что внутри (перенесено из привычек agat/train.py):
  * бюджет в эпохах (`--epochs`) или шагах (`--steps`), lr: прогрев → cosine → полка;
  * AdamW8bit (bitsandbytes), если он есть, иначе обычный AdamW — печатается, что взято;
  * дропаут подписи (`--drop-text`) для CFG и дропаут слов (`--drop-word`) внутри подписи;
  * сэмплы на старте (шаг 0), каждые `--sample-every`; держим `--keep` последних чекпойнтов и
    `--samples-keep` сеток; сохранение атомарное;
  * `--probe`: подбирает максимальный батч, который влезает в память, и печатает его;
  * EMA сглаживание печатаемого лосса (`--ema-decay`), `--seed` случайный по умолчанию (печатается);
  * метрики: mse_lat (общий, low t<0.3, high t>0.7), mse_pix, switch, tail(z0), hf(gen)/hf(AE), skip;
  * SIGUSR1 печатает стеки всех потоков (видно, ждёт процесс GPU или сеть, без py-spy).
"""
from __future__ import annotations

import argparse
import faulthandler
import json
import math
import os
import random
import shutil
import signal
import time
from pathlib import Path

os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")
signal.signal(signal.SIGINT, signal.default_int_handler)
faulthandler.register(signal.SIGUSR1)

import torch                      # noqa: E402
import torch.nn.functional as F   # noqa: E402

from model import DEV, UDiT        # noqa: E402
from text import EttinText         # noqa: E402
from ae.model import AutoEncoder, patchify   # noqa: E402
from ae_sdxs import SdxsAE                   # noqa: E402   (экспериментальная ветка: сторонний VAE)
from tfrecord_stream import TFRecordStream, count_shard_records   # noqa: E402  (шарды i1)


def load_ae(args):
    """Наш f8c32 AE или сторонний VAE (--vae sdxs): API совпадает, хвоста у второго нет."""
    if args.vae == "sdxs":
        return SdxsAE.from_pretrained(str(ROOT / "vae_sdxs"), DEV)
    return AutoEncoder.from_pretrained(str(ROOT / "vae"), device=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("--vae", choices=["f8c32", "sdxs"], default="f8c32",
                    help="f8c32 = наш vae/, sdxs = vae_sdxs/ (асимметричный декодер 2x, без хвоста)")
    ap.add_argument("--pix-weight", type=float, default=1.0,
                    help="вес пиксельного лосса через хвост; 0 = только латентный (для VAE без хвоста)")
    ap.add_argument("--graft", default="", help="привить блоки (стартуют тождеством), например 'e0=2,e1=2,d1=2'")
    ap.add_argument("--cfg", type=float, default=1.5, help="CFG при показе сэмплов")
    ap.add_argument("--sample-steps", type=int, default=16, help="шагов тела при сэмплинге")
    ap.add_argument("--tail-steps", type=int, default=8, help="шагов декодера AE (как в sample.py)")
    ap.add_argument("--lr-tail", type=float, default=3e-5)
    ap.add_argument("--tail", choices=["frozen", "train"], default="train")
    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("--seed", type=int, default=None)
    args = ap.parse_args()
    if args.vae == "sdxs":          # классическая схема: учим только латент, декодирует VAE
        args.pix_weight = 0.0
        args.tail = "frozen"
        print("[vae] sdxs: латентный режим — пиксельный лосс и хвост отключены", flush=True)

    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)

    has_tail = hasattr(ae, "decoder")      # у стороннего VAE (--vae sdxs) хвоста нет

    # ---------------------------------------------------------------- модель и оптимизатор
    m = UDiT(ae, text_dim=enc.dim).to(DEV)
    start = 0

    def load_body(sd):
        own = m.body.state_dict()
        take = {k: v for k, v in sd.items() if k in own and own[k].shape == v.shape}
        m.body.load_state_dict(take, strict=False)
        return len(take)

    if args.init:
        p = Path(args.init)
        if p.is_dir():
            from export import load_exported
            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"])
            if sd.get("decoder") is not None and args.tail == "train" and has_tail:
                ae.decoder.load_state_dict(sd["decoder"])
            start = int(sd.get("step", 0))
            print(f"[init] {p}: перенесено {n_take} тензоров, шаг {start}", flush=True)
    if args.resume and (out / "last.pt").exists():
        sd = torch.load(out / "last.pt", map_location=DEV)
        load_body(sd["body"]); enc.mixer.load_state_dict(sd["mixer"])
        if sd.get("decoder") is not None and args.tail == "train" and has_tail:
            ae.decoder.load_state_dict(sd["decoder"])
        start = int(sd.get("step", 0))
        print(f"[resume] {out / 'last.pt'} → шаг {start}", flush=True)

    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 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())
    n_tail = sum(p.numel() for p in ae.decoder.parameters()) if has_tail else 0
    print(f"[model] тело {n_body/1e6:.1f}M + миксер {sum(p.numel() for p in enc.trainable())/1e6:.2f}M "
          f"+ хвост {n_tail/1e6:.1f}M [{args.tail}]", flush=True)

    groups = [{"params": list(m.body.parameters()), "lr": args.lr},
              {"params": enc.trainable(), "lr": args.lr}]
    if args.tail == "train" and has_tail:
        ae.decoder.requires_grad_(True)
        groups.append({"params": list(ae.decoder.parameters()), "lr": args.lr_tail})
    params = [p for g in groups for p in g["params"]]
    try:
        import bitsandbytes as bnb
        opt = bnb.optim.AdamW8bit(groups, betas=(0.9, 0.95), weight_decay=0.0)
        print("[opt] AdamW8bit", flush=True)
    except Exception as e:
        opt = torch.optim.AdamW(groups, betas=(0.9, 0.95), weight_decay=0.0)
        print(f"[opt] AdamW fp32 (8bit недоступен: {type(e).__name__})", flush=True)

    lr_min = args.lr_min if args.lr_min is not None else args.lr * 0.1
    period = args.lr_period or steps              # длина cosine в шагах, отсчёт от start
    print(f"[lr] {args.lr:g}: прогрев {args.lr_warmup} от шага {start}, cosine до {lr_min:g} к шагу {start + period}", flush=True)

    # ---------------------------------------------------------------- показ сэмплов
    SAMPLE_PROMPTS = [
        "A photograph of a red fox sleeping on a mossy log in a misty forest at dawn.",
        "A city made of glass where rivers of light flow between the towers, cinematic lighting.",
        'A neon sign reading "hello world" mounted on a brick wall at night in the rain.',
        "An impressionist oil painting of a woman with a parasol standing by a lake.",
    ]
    show_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 = [], [], []
        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=args.tail_steps, cfg=args.cfg,
                         null_text=enc.mix(HS_E.float()), null_mask=MK_E,
                         generator=torch.Generator(DEV).manual_seed(0))
            tiles.append(Image.fromarray(((y[0].clamp(-1, 1) + 1) * 127.5).byte()
                                         .permute(1, 2, 0).cpu().numpy()))
            hfs.append(hf_energy(y)); ys.append(y)
        if show_x:
            ae_rec = ae.decode_raw(ae.encode(show_x[0]) * ae.latent_std + ae.latent_mean,
                                   steps=args.tail_steps, generator=torch.Generator(DEV).manual_seed(0))
            hf_ae = hf_energy(ae_rec)
        else:
            hf_ae = float("nan")
        W, H = tiles[0].size
        canvas = Image.new("RGB", (W * len(tiles), H), "white")
        for i, t in enumerate(tiles):
            canvas.paste(t, (i * W, 0))
        canvas.save(out / "samples" / "latest.png")
        canvas.save(out / "samples" / f"step_{step:07d}.png")
        rotate(out / "samples", args.samples_keep, "step_*.png")
        bright = float(torch.stack(ys).mean())
        contrast = float(torch.stack(ys).std())
        return hfs[0], hf_ae, canvas, bright, contrast

    def save(step):
        sd = {"body": {k: v.cpu() for k, v in m.body.state_dict().items()},
              "mixer": {k: v.cpu() for k, v in enc.mixer.state_dict().items()},
              "decoder": ({k: v.cpu() for k, v in ae.decoder.state_dict().items()}
                          if (args.tail == "train" and has_tail) else None),
              "step": step, "args": vars(args), "seed": seed}
        torch.save(sd, out / "last.pt")
        torch.save(sd, out / "ckpt" / f"step_{step:07d}.pt")
        rotate(out / "ckpt", args.keep, "step_*.pt")

    # ---------------------------------------------------------------- цикл
    lat_hw = (args.size // 8, args.size // 8)
    use_pix = bool(args.pix_weight > 0 and hasattr(m.ae, "_packing"))
    pkx = m.ae._packing([lat_hw] * args.batch) if use_pix else None
    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()
                if use_pix:
                    raw = m.to_raw(z0 + (z_hat - z0).detach())
                    tx = torch.rand(args.batch, device=DEV)
                    ex = torch.randn_like(x0)
                    xt = (1 - tx[:, None, None, None]) * x0 + tx[:, None, None, None] * ex
                    xt_packed = pkx.pack([patchify(xt[i].to(ae.compute_dtype), 8) for i in range(args.batch)])
                    z_packed = pkx.pack([raw[i].flatten(1).transpose(0, 1) for i in range(args.batch)])
                    preds = ae.decoder(xt_packed, tx, z_packed, pkx)
                    pix_terms = torch.stack([(p.float() - x0[i]).pow(2).mean() for i, p in enumerate(preds)])
                    loss_p = pix_terms.mean()
                else:
                    loss_p = torch.zeros((), device=DEV)
            total = loss_l + args.pix_weight * loss_p
            if not torch.isfinite(total):
                skips += 1
                continue
            (total / args.accum).backward()
            ok += 1
            run["lat"] += loss_l.item(); run["pix"] += loss_p.item()
            if low.any():
                run["low"] += per_elem[low].mean().item(); run["n_low"] += 1
            if high.any():
                run["high"] += per_elem[high].mean().item(); run["n_high"] += 1
        if not ok:
            print(f"  ! шаг {step}: не-конечный лосс, пропуск (всего {skips})", flush=True)
            continue
        run["n"] += 1
        torch.nn.utils.clip_grad_norm_(params, 1.0)
        lr = lr_at(step, args.lr, lr_min, args.lr_warmup, period, start)
        for g, base in zip(opt.param_groups, (args.lr, args.lr, args.lr_tail)):
            g["lr"] = lr * (base / args.lr)
        opt.step()

        if step % args.log_every == 0:
            d = args.ema_decay
            cur = {"lat": run["lat"] / run["n"], "pix": run["pix"] / run["n"],
                   "low": run["low"] / max(run["n_low"], 1), "high": run["high"] / max(run["n_high"], 1)}
            for k, v in cur.items():
                ema[k] = v if k not in ema else d * ema[k] + (1 - d) * v
            el = time.perf_counter() - t0; inc = el - prev_el; prev_el = el
            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}) | mse_pix {cur['pix']:.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:
            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")          # метрика хвоста; у стороннего VAE его нет
            if show_x and has_tail:           # и размер декода должен совпасть с кадром (sdxs отдаёт 2x)
                z_ref = ae.encode(show_x[0])
                rec = ae.decode_raw(m.to_raw(z_ref), steps=args.tail_steps,
                                    generator=torch.Generator(DEV).manual_seed(0))
                if rec.shape == show_x[0].shape:
                    tail_psnr = psnr(rec, show_x[0])
            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"{ema.get('pix', 0):.5f},{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")
        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()