File size: 4,372 Bytes
b4f5f26 f3de28c 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 | #!/usr/bin/env python3
"""Сборка публикационного дерева: код + доки + веса в fp16.
python3 src/publish.py --out tmp/publish # только собрать
python3 src/publish.py --out tmp/publish --push # собрать и залить (REPO_ID)
Локальные fp32-веса остаются на месте: они нужны обучению, а на Hub едет fp16.
"""
from __future__ import annotations
import argparse
import json
import shutil
from pathlib import Path
import torch
from safetensors.torch import load_file, save_file
ROOT = Path(__file__).resolve().parent.parent
REPO_ID = "AiArtLab/Ra"
COPY = ["README.md", ".gitattributes", ".gitignore", "docs", "src",
"vae/config.json", "vae/LICENSE", "vae/README.md", "vae/.gitattributes", "vae/example.py",
"vae/ae", "text_encoder/config.json", "text_encoder/tokenizer.json",
"text_encoder/tokenizer_config.json", "text_encoder/special_tokens_map.json",
"text_encoder/README.md", "text_encoder/.gitattributes",
"transformer/config.json"]
def half_file(src: Path, dst: Path, key_prefix: str = ""):
tensors = load_file(str(src))
out, overflow = {}, 0
for k, v in tensors.items():
name = f"{key_prefix}{k}" if key_prefix else k
if v.is_floating_point():
if v.numel() and float(v.abs().max()) > 65504:
overflow += 1
v = v.half()
out[name] = v.contiguous()
assert torch.isfinite(out[name]).all(), f"{name}: не-конечные значения после fp16"
dst.parent.mkdir(parents=True, exist_ok=True)
save_file(out, str(dst))
print(f"[fp16] {dst.relative_to(dst.parents[1])}: {len(out)} тензоров, "
f"{src.stat().st_size / 2**20:.0f} → {dst.stat().st_size / 2**20:.0f} МиБ"
+ (f", ВНИМАНИЕ: {overflow} тензоров вне диапазона fp16" if overflow else ""), flush=True)
def main():
ap = argparse.ArgumentParser()
ap.add_argument("--out", default=str(ROOT / "tmp" / "publish"))
ap.add_argument("--repo", default=REPO_ID)
ap.add_argument("--push", action="store_true", help="залить собранное на Hub (публичный репозиторий)")
args = ap.parse_args()
out = Path(args.out)
if out.exists():
shutil.rmtree(out)
out.mkdir(parents=True)
for rel in COPY:
src = ROOT / rel
dst = out / rel
dst.parent.mkdir(parents=True, exist_ok=True)
if src.is_dir():
shutil.copytree(src, dst, ignore=shutil.ignore_patterns("__pycache__", ".ipynb_checkpoints"))
else:
shutil.copy2(src, dst)
half_file(ROOT / "transformer" / "diffusion_pytorch_model.safetensors",
out / "transformer" / "diffusion_pytorch_model.safetensors")
half_file(ROOT / "vae" / "model.safetensors", out / "vae" / "model.safetensors")
sd = torch.load(ROOT / "text_encoder" / "pytorch_model.bin", map_location="cpu", weights_only=True)
sd = {k: v for k, v in sd.items() if not k.startswith(("head.", "decoder."))} # MLM-голова не нужна
save_file({k: v.half().contiguous() for k, v in sd.items()},
str(out / "text_encoder" / "model.safetensors"))
print(f"[fp16] text_encoder/model.safetensors: {len(sd)} тензоров "
f"({sum(v.numel() for v in sd.values()) / 1e6:.1f}M), "
f"{(out / 'text_encoder' / 'model.safetensors').stat().st_size / 2**20:.0f} МиБ", flush=True)
cfg_path = out / "transformer" / "config.json"
cfg = json.loads(cfg_path.read_text())
cfg["storage_dtype"] = "float16"
cfg_path.write_text(json.dumps(cfg, indent=2, ensure_ascii=False))
total = sum(f.stat().st_size for f in out.rglob("*") if f.is_file())
print(f"[готово] {out}: {total / 2**30:.2f} ГиБ", flush=True)
if args.push:
from huggingface_hub import HfApi
api = HfApi()
api.create_repo(args.repo, repo_type="model", private=False, exist_ok=True)
api.upload_folder(folder_path=str(out), repo_id=args.repo, repo_type="model",
commit_message=f"Ra fp16 (step {cfg['step']})")
print(f"[push] https://huggingface.co/{args.repo}", flush=True)
if __name__ == "__main__":
main()
|