Blame for src/publish.py

#!/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", "requirements.txt", "run_sdxs.sh", ".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("--vae-dir", default="vae", help="папка VAE (vae или vae_sdxs)")
    ap.add_argument("--transformer-dir", default="transformer")
    ap.add_argument("--readme", default="README.md", help="какой README положить в корень")
    ap.add_argument("--push", action="store_true", help="залить собранное на Hub (публичный репозиторий)")
    ap.add_argument("--step", type=int, default=None, help="шаг чекпойнта — только для сообщения коммита")
    args = ap.parse_args()
    out = Path(args.out)
    if out.exists():
        shutil.rmtree(out)
    out.mkdir(parents=True)

    rels = list(COPY) + [args.readme, f"{args.vae_dir}/config.json", f"{args.vae_dir}/README.md"]
    if args.transformer_dir != "transformer":
        rels.append(f"{args.transformer_dir}/config.json")
    for rel in rels:
        src = ROOT / rel
        if not src.exists():
            continue
        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:
            dst_rel = "README.md" if rel == args.readme else rel
            shutil.copy2(src, (out / dst_rel))

    tdir, vdir = ROOT / args.transformer_dir, ROOT / args.vae_dir
    half_file(tdir / "diffusion_pytorch_model.safetensors",
              out / args.transformer_dir / "diffusion_pytorch_model.safetensors")
    vae_weights = vdir / "model.safetensors"
    if not vae_weights.exists():                       # AsymmetricAutoencoderKL кладёт так
        vae_weights = vdir / "diffusion_pytorch_model.safetensors"
    half_file(vae_weights, out / args.vae_dir / vae_weights.name)

    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 / args.transformer_dir / "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 CommitOperationAdd, HfApi
        api = HfApi()
        api.create_repo(args.repo, repo_type="model", private=False, exist_ok=True)
        # create_commit вместо upload_folder: upload_folder уважает .gitignore из бандла
        # (а там `vae_sdxs/*.safetensors`) и молча не заливает веса
        ops = [CommitOperationAdd(path_in_repo=p.relative_to(out).as_posix(), path_or_fileobj=str(p))
               for p in sorted(out.rglob("*")) if p.is_file()]
        api.create_commit(repo_id=args.repo, repo_type="model", operations=ops,
                          commit_message=(f"Ra sdxs-VAE (step {args.step})" if args.step
                                          else "Ra sdxs-VAE (weights)"))
        print(f"[push] https://huggingface.co/{args.repo}: {len(ops)} файлов", flush=True)


if __name__ == "__main__":
    main()