96 lines
4.1 KiB
Python
96 lines
4.1 KiB
Python
# /// script
|
|
# requires-python = ">=3.11,<3.12"
|
|
# dependencies = ["torch==2.13.0", "onnx==1.19.0", "numpy==1.26.4", "einops==0.8.2", "gdown==6.2.0"]
|
|
# [tool.uv.sources]
|
|
# torch = { index = "pytorch" }
|
|
# [[tool.uv.index]]
|
|
# name = "pytorch"
|
|
# url = "https://download.pytorch.org/whl/cpu"
|
|
# explicit = true
|
|
# ///
|
|
"""Export the official HAT Sharper checkpoint with FP16 weight storage."""
|
|
|
|
import argparse
|
|
import hashlib
|
|
import json
|
|
from pathlib import Path
|
|
from urllib.request import urlopen
|
|
|
|
import gdown
|
|
import onnx
|
|
import torch
|
|
|
|
from onnx_weights import compress_weights
|
|
|
|
|
|
REVISION = "1638a9a822581657811867bf670717f8371fc3e5"
|
|
CHECKPOINT_SHA256 = "5800b67136006eb8cab3b4ed7c8d73b6a195bb18e6cc709b674f9aa069c00271"
|
|
SOURCE_URL = f"https://raw.githubusercontent.com/XPixelGroup/HAT/{REVISION}"
|
|
|
|
|
|
def load_model(checkpoint, cache):
|
|
architecture = cache / "hat_arch.py"
|
|
if not architecture.exists():
|
|
architecture.write_bytes(urlopen(SOURCE_URL + "/hat/archs/hat_arch.py", timeout=60).read())
|
|
if hashlib.sha256(architecture.read_bytes()).hexdigest() != "81d8cecf491975246c9ebb20480898c656f59de020f38b05daa68741e426117f":
|
|
raise RuntimeError("HAT architecture SHA-256 mismatch")
|
|
# Export only: remove BasicSR's training registry and use the equivalent
|
|
# PyTorch initialization utilities. The official forward remains unchanged.
|
|
source = architecture.read_text(encoding="utf-8")
|
|
source = source.replace("from basicsr.utils.registry import ARCH_REGISTRY", "")
|
|
source = source.replace("@ARCH_REGISTRY.register()", "")
|
|
source = source.replace(
|
|
"from basicsr.archs.arch_util import to_2tuple, trunc_normal_",
|
|
"from torch.nn.modules.utils import _pair as to_2tuple\nfrom torch.nn.init import trunc_normal_",
|
|
)
|
|
namespace = {"__name__": "hat_arch"}
|
|
exec(compile(source, str(architecture), "exec"), namespace)
|
|
model = namespace["HAT"](
|
|
upscale=4, in_chans=3, img_size=64, window_size=16,
|
|
compress_ratio=3, squeeze_factor=30, conv_scale=0.01,
|
|
overlap_ratio=0.5, img_range=1., depths=[6] * 6,
|
|
embed_dim=180, num_heads=[6] * 6, mlp_ratio=2,
|
|
upsampler="pixelshuffle", resi_connection="1conv",
|
|
).eval()
|
|
with checkpoint.open("rb") as stream:
|
|
if hashlib.file_digest(stream, "sha256").hexdigest() != CHECKPOINT_SHA256:
|
|
raise RuntimeError("HAT Sharper checkpoint SHA-256 mismatch")
|
|
model.load_state_dict(torch.load(checkpoint, map_location="cpu", weights_only=True)["params_ema"], strict=True)
|
|
return model
|
|
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser(description=__doc__)
|
|
parser.add_argument("--destination", type=Path, required=True)
|
|
parser.add_argument("--checkpoint", type=Path)
|
|
args = parser.parse_args()
|
|
cache = Path(__file__).resolve().parents[2] / ".model-cache" / "hat-sharper"
|
|
cache.mkdir(parents=True, exist_ok=True)
|
|
checkpoint = args.checkpoint or cache / "Real_HAT_GAN_sharper.pth"
|
|
if not checkpoint.exists():
|
|
gdown.download(id="1EioFq5-mKmv1uqta_Byd9cgXp9SU3zjj", output=str(checkpoint))
|
|
torch.set_num_threads(8)
|
|
torch.manual_seed(0)
|
|
model = load_model(checkpoint, cache)
|
|
args.destination.mkdir(parents=True, exist_ok=True)
|
|
output = args.destination / "model.onnx"
|
|
with torch.inference_mode():
|
|
torch.onnx.export(model, torch.zeros(1, 3, 256, 256), str(output),
|
|
opset_version=17, input_names=["input"], output_names=["output"],
|
|
do_constant_folding=False, dynamo=False)
|
|
graph = onnx.load(output)
|
|
compress_weights(graph)
|
|
onnx.checker.check_model(graph)
|
|
onnx.save(graph, output)
|
|
(args.destination / "LICENSE.txt").write_bytes(urlopen(SOURCE_URL + "/LICENSE", timeout=60).read())
|
|
if sum(path.stat().st_size for path in args.destination.iterdir()) >= 60_000_000:
|
|
raise RuntimeError("HAT asset exceeds 60 MB")
|
|
with output.open("rb") as stream:
|
|
digest = hashlib.file_digest(stream, "sha256").hexdigest()
|
|
print(json.dumps({"bytes": output.stat().st_size, "sha256": digest,
|
|
"revision": REVISION, "checkpoint_sha256": CHECKPOINT_SHA256}), flush=True)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|