| |
| """Сборка публикационного дерева: код + доки + веса в 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(): |
| 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."))} |
| 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) |
| |
| |
| 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() |
|
|
