| """Экспорт обученного UDM в `transformer/`: safetensors + config.json. |
| |
| Три части в одном файле, по префиксам: |
| body.* — тело (8 блоков из энкодера AE + 4 новых), поток в латентном пространстве |
| text_mixer.* — обучаемый выбор слоёв Ettin (SliceMixer) |
| tail.* — декодер AE (рендер латента в пиксели), дообученный как обычный flow |
| |
| python3 src/export.py --ckpt tmp/run256/last.pt --out transformer |
| """ |
| from __future__ import annotations |
|
|
| import argparse |
| import json |
| from pathlib import Path |
|
|
| import torch |
| from safetensors.torch import save_file |
|
|
| ROOT = Path(__file__).resolve().parent.parent |
|
|
|
|
| def main(): |
| ap = argparse.ArgumentParser() |
| ap.add_argument("--ckpt", default=str(ROOT / "tmp" / "run256" / "last.pt")) |
| ap.add_argument("--out", default=str(ROOT / "transformer")) |
| ap.add_argument("--metrics", default="", help="дописать метрики в config.json (строка)") |
| ap.add_argument("--half", action="store_true", help="хранить веса в fp16 (вдвое меньше, для публикации)") |
| args = ap.parse_args() |
|
|
| ck = torch.load(args.ckpt, map_location="cpu") |
| tensors = {} |
| for name, sd in (("body", ck["body"]), ("text_mixer", ck["mixer"]), ("tail", ck.get("decoder") or {})): |
| for k, v in sd.items(): |
| tensors[f"{name}.{k}"] = (v.half() if args.half else v).contiguous() |
| out = Path(args.out) |
| out.mkdir(parents=True, exist_ok=True) |
| save_file(tensors, str(out / "diffusion_pytorch_model.safetensors")) |
|
|
| ae_cfg = json.loads((ROOT / "vae" / "config.json").read_text()) |
| ta = ck.get("args") or {} |
| counts = {} |
| for k in ck["body"]: |
| g = k.split(".")[0] |
| if g[0] in "ed" and "." in k: |
| idx = k.split(".")[1] |
| if idx.isdigit(): |
| counts.setdefault(g, set()).add(idx) |
| nblk = {g: len(v) for g, v in counts.items()} |
| |
| cfg = { |
| "_class_name": "UDiT", |
| "format": "udm/v1", |
| "storage_dtype": "float16" if args.half else "float32", |
| "latent_channels": ae_cfg["latent_channels"], |
| "text_dim": ta.get("text_dim", 768), |
| "body": { |
| "kind": "пирамидальный U: ширина растёт к низу", |
| "widths": [192, 512, 1024, 1536], |
| "heads": [6, 8, 8, 8], |
| "blocks": nblk, |
| "levels": [ |
| {"level": 0, "tokens": "÷1", "width": 192, |
| "blocks": "conv 6 enc / 4 dec + 2 self-attn на сторону"}, |
| {"level": 1, "tokens": "÷2", "width": 512, "blocks": "attn 5 enc / 4 dec"}, |
| {"level": 2, "tokens": "÷4", "width": 1024, "blocks": "attn 5 enc / 4 dec"}, |
| {"level": 3, "tokens": "÷8", "width": 1536, "blocks": "attn 4 (боттлнек)"}, |
| ], |
| "skips": "аддитивные (SD-стиль): e0/e1/e2 прибавляются в декодере", |
| "down": "mean-pool 2x + Linear", |
| "up": "Linear + nearest 2x", |
| "positions": { |
| "rope": {"kind": "2d", "base": 100, "phase": "в латентных пикселях на каждом уровне"}, |
| "absolute": {"kind": "2d sincos", "dim": 64, |
| "into": "zero-init Linear(64 -> width) к выходу stem, в паре с RoPE"}, |
| }, |
| "text": "кросс-внимание на всех уровнях", |
| "parametrization": "x0 латента", |
| "params": round(sum(v.numel() for k, v in tensors.items() if k.startswith("body.")) / 1e6, 2), |
| }, |
| } |
| (out / "config.json").write_text(json.dumps(cfg, indent=2, ensure_ascii=False)) |
|
|
| print(f"сохранено: {out}/diffusion_pytorch_model.safetensors " |
| f"({sum(v.numel() for v in tensors.values())/1e6:.1f}M параметров, " |
| f"{(out / 'diffusion_pytorch_model.safetensors').stat().st_size/1e6:.0f} МБ)") |
| print(f"тело {cfg['body']['params']}M, блоки {nblk}") |
|
|
|
|
| def load_exported(model, enc, ae, path, device=None): |
| """Загрузить из `transformer/` (config.json + safetensors с префиксами body./text_mixer./tail.).""" |
| import json |
| from safetensors.torch import load_file |
| path = Path(path) |
| sd = load_file(str(path / "diffusion_pytorch_model.safetensors"), device=str(device) if device else None) |
| def sub(prefix): |
| return {k[len(prefix) + 1:]: v for k, v in sd.items() if k.startswith(prefix + ".")} |
| model.body.load_state_dict(sub("body"), strict=False) |
| enc.mixer.load_state_dict(sub("text_mixer")) |
| if getattr(ae, "decoder", None) is not None and any(k.startswith("tail.") for k in sd): |
| ae.decoder.load_state_dict(sub("tail")) |
| cfg_path = path / "config.json" |
| return json.loads(cfg_path.read_text()) if cfg_path.exists() else {} |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
