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