File size: 5,438 Bytes
b4f5f26 cb6884f 9b3ae10 cb6884f b4f5f26 cb6884f b4f5f26 cb6884f 9b3ae10 b4f5f26 cb6884f b4f5f26 cb6884f 234d8d3 cb6884f b4f5f26 cb6884f b4f5f26 cb6884f b4f5f26 8096951 b4f5f26 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 | """Экспорт обученного 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()}
# конфиг описывает ТОЛЬКО тело; статистики латента — в vae/config.json, шаг — параметр запуска
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")) # хвоста может не быть (--vae sdxs)
cfg_path = path / "config.json"
return json.loads(cfg_path.read_text()) if cfg_path.exists() else {}
if __name__ == "__main__":
main()
|