| """Экспорт обученного 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 {} |
| training = { |
| "dataset": ta.get("hf_dataset") or ta.get("data") or "~/dataset/butterflies_256x64", |
| "parts": ta.get("parts"), |
| "size": ta.get("size"), |
| "steps": int(ck["step"]), |
| "batch": ta.get("batch", 8), |
| "accum": ta.get("accum", 1), |
| "lr": ta.get("lr", 1e-4), |
| "lr_tail": ta.get("lr_tail", 3e-5), |
| "drop_text": ta.get("drop_text", 0.1), |
| "init": ta.get("init"), |
| "metrics": args.metrics, |
| } |
| cfg = { |
| "_class_name": "UDiT", |
| "name": "Ra", |
| "format": "udm/v1", |
| "step": int(ck["step"]), |
| "latent": { |
| "channels": ae_cfg["latent_channels"], |
| "spatial_factor": ae_cfg["patch_size"], |
| "mean": ae_cfg["latent_mean"], |
| "std": ae_cfg["latent_std"], |
| "source": "vae/ (Muinez/f8c32-ae)", |
| }, |
| "text": { |
| "encoder": "text_encoder/ (jhu-clsp/ettin-encoder-150m, 768-dim, 22 слоя)", |
| "max_len": 192, |
| "slices": [2, 5, 8, 11, 14, 17, 20], |
| "slice_blocks": 2, |
| "trainable": "text_mixer.*", |
| }, |
| "body": { |
| "kind": "пирамидальный U: ширина растёт к низу", |
| "levels": [ |
| {"level": 0, "tokens": "÷1", "width": 192, "blocks": "conv 6 enc / 4 dec + 1 self-attn на сторону", |
| "text": "кросс-внимание в каждом conv-блоке и в self-attn"}, |
| {"level": 1, "tokens": "÷2", "width": 512, "blocks": "attn 4 enc / 3 dec", "text": "кросс-внимание"}, |
| {"level": 2, "tokens": "÷4", "width": 1024, "blocks": "attn 3 enc / 2 dec", "text": "кросс-внимание"}, |
| {"level": 3, "tokens": "÷8", "width": 1536, "blocks": "attn 2 (боттлнек)", "text": "кросс-внимание"}, |
| ], |
| "skips": "аддитивные (SD-стиль): e0/e1/e2 прибавляются в декодере", |
| "down": "mean-pool 2× + Linear", |
| "up": "Linear + nearest 2×", |
| "rope": "2D RoPE, base 100, фаза в латентных пикселях на каждом уровне", |
| "parametrization": "x0 латента", |
| }, |
| "tail": { |
| "source": "декодер AE, 6 блоков @512, PixelShuffle ×3", |
| "trained_as": "flow по пикселям, условие z0 + (ошибка тела).detach()", |
| "sample_steps": 4, |
| }, |
| "params": { |
| "body": round(sum(v.numel() for k, v in tensors.items() if k.startswith("body.")) / 1e6, 2), |
| "text_mixer": round(sum(v.numel() for k, v in tensors.items() |
| if k.startswith("text_mixer.")) / 1e6, 2), |
| "tail": round(sum(v.numel() for k, v in tensors.items() if k.startswith("tail.")) / 1e6, 2), |
| }, |
| "training": {k: v for k, v in training.items() if v is not None}, |
| "storage_dtype": "float16" if args.half else "float32", |
| } |
| cfg["params"]["total"] = round(sum(cfg["params"].values()), 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['params']}") |
|
|
|
|
| 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")) |
| 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() |
|
|
