|
|
| """导出 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()
|
|
|
| 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()
|
|
|
