110 lines
3.5 KiB
Python
110 lines
3.5 KiB
Python
"""Export official PyTorch super-resolution models to ONNX."""
|
|
|
|
import numpy as np
|
|
|
|
|
|
OPSET_VERSION = 17
|
|
|
|
|
|
def _require_build_dependencies():
|
|
try:
|
|
import onnx
|
|
import onnxruntime
|
|
import torch
|
|
except ImportError as error:
|
|
raise RuntimeError(
|
|
"Model build dependencies are missing. Run with: "
|
|
"uv sync --group models"
|
|
) from error
|
|
return onnx, onnxruntime, torch
|
|
|
|
|
|
def _load_weights(torch, checkpoint_path, key):
|
|
checkpoint = torch.load(checkpoint_path, map_location="cpu", weights_only=True)
|
|
return checkpoint[key]
|
|
|
|
|
|
def _validate_export(onnx, onnxruntime, torch, model, sample, destination):
|
|
onnx.checker.check_model(onnx.load(destination))
|
|
expected = model(sample).detach().numpy()
|
|
session = onnxruntime.InferenceSession(
|
|
str(destination),
|
|
providers=["CPUExecutionProvider"],
|
|
)
|
|
actual = session.run(None, {"input": sample.numpy()})[0]
|
|
np.testing.assert_allclose(actual, expected, rtol=1e-4, atol=1e-5)
|
|
|
|
|
|
def _export(torch, model, sample, destination, dynamic_axes=None):
|
|
model.eval()
|
|
torch.onnx.export(
|
|
model,
|
|
sample,
|
|
destination,
|
|
export_params=True,
|
|
opset_version=OPSET_VERSION,
|
|
do_constant_folding=True,
|
|
input_names=["input"],
|
|
output_names=["output"],
|
|
dynamic_axes=dynamic_axes,
|
|
dynamo=False,
|
|
)
|
|
|
|
|
|
def _restore_realesrgan_scope_names(onnx, destination):
|
|
model = onnx.load(destination)
|
|
for node in model.graph.node:
|
|
node.name = node.name.replace("/body.", "/body/body.")
|
|
for values in (node.input, node.output):
|
|
for index, value in enumerate(values):
|
|
values[index] = value.replace("/body.", "/body/body.")
|
|
onnx.save(model, destination)
|
|
|
|
|
|
def export_realesrgan_wdn_x4(checkpoint_path, destination):
|
|
onnx, onnxruntime, torch = _require_build_dependencies()
|
|
nn = torch.nn
|
|
|
|
class SRVGGNetCompact(nn.Module):
|
|
def __init__(self):
|
|
super().__init__()
|
|
self.upscale = 4
|
|
self.body = nn.ModuleList([nn.Conv2d(3, 64, 3, 1, 1), nn.PReLU(64)])
|
|
for _ in range(32):
|
|
self.body.extend([nn.Conv2d(64, 64, 3, 1, 1), nn.PReLU(64)])
|
|
self.body.append(nn.Conv2d(64, 3 * self.upscale**2, 3, 1, 1))
|
|
self.upsampler = nn.PixelShuffle(self.upscale)
|
|
|
|
def forward(self, value):
|
|
result = value
|
|
for layer in self.body:
|
|
result = layer(result)
|
|
result = self.upsampler(result)
|
|
return result + torch.nn.functional.interpolate(
|
|
value,
|
|
scale_factor=self.upscale,
|
|
mode="nearest",
|
|
)
|
|
|
|
model = SRVGGNetCompact()
|
|
model.load_state_dict(_load_weights(torch, checkpoint_path, "params"), strict=True)
|
|
sample = torch.rand(1, 3, 16, 16)
|
|
dynamic_axes = {
|
|
"input": {2: "height", 3: "width"},
|
|
"output": {2: "height4", 3: "width4"},
|
|
}
|
|
_export(torch, model, sample, destination, dynamic_axes)
|
|
_restore_realesrgan_scope_names(onnx, destination)
|
|
_validate_export(onnx, onnxruntime, torch, model, sample, destination)
|
|
|
|
|
|
def export_model(exporter, checkpoint_path, destination):
|
|
exporters = {
|
|
"realesrgan_wdn_x4": export_realesrgan_wdn_x4,
|
|
}
|
|
try:
|
|
export = exporters[exporter]
|
|
except KeyError as error:
|
|
raise ValueError(f"Unknown model exporter: {exporter}") from error
|
|
export(checkpoint_path, destination)
|