File size: 3,957 Bytes
9f08d74 | 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 | """Official PNG codec compatible weight reconstruction adapters."""
from __future__ import annotations
import hashlib
import json
from collections.abc import Mapping, Sequence
from pathlib import Path
import numpy as np
from PIL import Image
def decode_weight_image(image_path: str | Path, manifest_path: str | Path) -> tuple[np.ndarray, list[str]]:
"""公式manifest順でRGBのR/Gをfp16 bit patternとしてfloat32へ復元する。"""
metadata = json.loads(Path(manifest_path).read_text())
pixels = np.asarray(Image.open(image_path).convert("RGB"), dtype=np.uint8).reshape(-1, 3)
total = int(metadata["total_parameters"])
if len(pixels) < total:
raise ValueError(f"weight image has {len(pixels)} pixels, needs {total}")
bits = (pixels[:total, 0].astype(np.uint16) << 8) | pixels[:total, 1].astype(np.uint16)
return bits.view(np.float16).astype(np.float32), [str(item["name"]) for item in metadata["params"]]
def flatten_manifest(arrays: Mapping[str, np.ndarray]) -> tuple[np.ndarray, list[tuple[str, int, int]]]:
"""mappingの挿入順を保持し、float16 quantization後のfloat32 flatを返す。"""
parts: list[np.ndarray] = []
manifest: list[tuple[str, int, int]] = []
offset = 0
for name, value in arrays.items():
part = np.asarray(value).reshape(-1).astype(np.float16).astype(np.float32)
parts.append(part)
manifest.append((name, offset, offset + len(part)))
offset += len(part)
return np.concatenate(parts) if parts else np.empty(0, dtype=np.float32), manifest
def layer_offsets(names: Sequence[str], sizes: Sequence[int]) -> dict[str, tuple[int, int]]:
"""layer名とnumelからmetric用offsetを構築する。"""
offsets: dict[str, tuple[int, int]] = {}
offset = 0
for name, size in zip(names, sizes):
offsets[name] = (offset, offset + int(size))
offset += int(size)
return offsets
def sha256_file(path: str | Path) -> str:
"""ファイルのSHA-256を返す。"""
digest = hashlib.sha256()
with Path(path).open("rb") as stream:
for chunk in iter(lambda: stream.read(1024 * 1024), b""):
digest.update(chunk)
return digest.hexdigest()
def reconstruct_dit(image_path: str | Path, manifest_path: str | Path, device: str = "cpu", config_path: str | Path | None = None):
"""model_png.jsonのname/shape順でDiTを再構築し、copy後の全parameterをfinite gateする。"""
import torch
metadata = json.loads(Path(manifest_path).read_text(encoding="utf-8"))
config = json.loads(Path(config_path).read_text(encoding="utf-8"))["dit"] if config_path is not None else metadata["cfg"]
from dit import DiT
model = DiT(dim=int(config["dim"]), depth=int(config["depth"]), heads=int(config["heads"]))
named = dict(model.named_parameters())
pixels = np.asarray(Image.open(image_path).convert("RGB"), dtype=np.uint8).reshape(-1, 3)
total = int(metadata["total_parameters"])
if pixels.shape[0] < total:
raise ValueError("weight image is shorter than manifest")
bits = (pixels[:total, 0].astype(np.uint16) << 8) | pixels[:total, 1].astype(np.uint16)
flat = bits.view(np.float16)
offset = 0
with torch.no_grad():
for item in metadata["params"]:
name = str(item["name"])
shape = tuple(int(value) for value in item["shape"])
numel = int(item["numel"])
if name not in named or tuple(named[name].shape) != shape or numel != int(np.prod(shape)):
raise ValueError(f"manifest mismatch for {name}")
named[name].copy_(torch.from_numpy(flat[offset:offset + numel].copy()).view(shape).float())
offset += numel
if offset != total or any(not bool(torch.isfinite(parameter).all()) for parameter in model.parameters()):
raise ValueError("finite gate failed after model parameter copy")
return model.to(device).eval()
|