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