File size: 2,823 Bytes
4811c23 | 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 | #!/usr/bin/env python
"""导出 torch.jit(fp16) 512x512 模型: 输入 LR[1,3,512,512][-1,1] -> 输出同尺寸。
内部: bicubic 512->128 -> 官方 4x 学生全链 -> 512。 (AdaIN 后处理不进测速模型)
用法: python src/export_jit.py --net weight/s2/net_params_X.pkl --out model_dir
"""
import argparse, os, sys
from pathlib import Path
REPO = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(REPO)); sys.path.insert(0, str(REPO / "src")); sys.path.insert(0, str(REPO / "official"))
import torch, torch.nn as nn, torch.nn.functional as F
from common import load_diffusers_sd, load_pruned_decoder, build_net
class SR512(nn.Module):
def __init__(self, net, tail):
super().__init__()
self.net = net
self.tail = tail
def forward(self, x512):
x128 = F.interpolate(x512, size=(128, 128), mode="bicubic", align_corners=False)
z = self.net(x128)
return self.tail(z)
def main():
ap = argparse.ArgumentParser()
ap.add_argument("--net", required=True)
ap.add_argument("--out", default="model_dir")
ap.add_argument("--half_decoder", default="weight/pretrained/halfDecoder.ckpt")
ap.add_argument("--model_id", default="models/stable-diffusion-2-1-base")
ap.add_argument("--name", default="your_model.pt")
args = ap.parse_args()
os.makedirs(args.out, exist_ok=True)
device = "cuda" if torch.cuda.is_available() else "cpu"
vae, unet, _, _ = load_diffusers_sd(args.model_id, dtype=torch.float32, device="cpu")
del vae
decoder = load_pruned_decoder(args.half_decoder, device="cpu", dtype=torch.float32)
net = build_net(unet, decoder)
sd = torch.load(args.net, map_location="cpu", weights_only=False)
if any(k.startswith("module.") for k in sd):
sd = {k.replace("module.", "", 1): v for k, v in sd.items()}
net.load_state_dict(sd, strict=True)
net.eval()
tail = nn.Sequential(*decoder.up_blocks, decoder.conv_norm_out,
decoder.conv_act, decoder.conv_out).eval()
model = SR512(net, tail).to(device).half().eval()
# trace on fixed 512 fp16
dummy = torch.randn(1, 3, 512, 512, device=device).half() * 0.5
with torch.no_grad():
traced = torch.jit.trace(model, dummy, check_trace=False)
traced = torch.jit.freeze(traced)
out_path = os.path.join(args.out, args.name)
traced.save(out_path)
# 自检: 两次前向一致性 + 形状
with torch.no_grad():
o1 = traced(dummy); o2 = traced(dummy)
assert o1.shape == dummy.shape, o1.shape
err = (o1 - o2).abs().max().item()
print("saved", out_path, "| deterministic max-diff:", err)
print("output range sample:", float(o1.min()), float(o1.max()))
if __name__ == "__main__":
main()
|