nan
89087fc2e2
- Implemented the "O Image Relief Plane" Geometry Nodes group to create a fixed-base image relief using depth images and various parameters. - Created the "O Shadeless" shader node group to display color without scene lighting while preserving shadows and reflections. - Introduced a preview feature for node groups, allowing for interactive visualization of node arrangements in an HTML format. - Added a template HTML file for rendering node previews with customizable views and controls.
205 lines
6.8 KiB
Python
205 lines
6.8 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"
|
|
) 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_safmn_real_x4(checkpoint_path, destination):
|
|
onnx, onnxruntime, torch = _require_build_dependencies()
|
|
nn = torch.nn
|
|
functional = torch.nn.functional
|
|
|
|
class LayerNorm(nn.Module):
|
|
def __init__(self, channels):
|
|
super().__init__()
|
|
self.weight = nn.Parameter(torch.ones(channels))
|
|
self.bias = nn.Parameter(torch.zeros(channels))
|
|
self.eps = 1e-6
|
|
|
|
def forward(self, value):
|
|
mean = value.mean(1, keepdim=True)
|
|
variance = (value - mean).pow(2).mean(1, keepdim=True)
|
|
value = (value - mean) / torch.sqrt(variance + self.eps)
|
|
return self.weight[:, None, None] * value + self.bias[:, None, None]
|
|
|
|
class SAFM(nn.Module):
|
|
def __init__(self, channels):
|
|
super().__init__()
|
|
chunk_channels = channels // 4
|
|
self.mfr = nn.ModuleList(
|
|
nn.Conv2d(chunk_channels, chunk_channels, 3, 1, 1, groups=chunk_channels)
|
|
for _ in range(4)
|
|
)
|
|
self.aggr = nn.Conv2d(channels, channels, 1)
|
|
self.act = nn.GELU()
|
|
|
|
def forward(self, value):
|
|
height, width = value.shape[-2:]
|
|
outputs = []
|
|
for level, chunk in enumerate(value.chunk(4, dim=1)):
|
|
if level:
|
|
size = (height // 2**level, width // 2**level)
|
|
chunk = functional.adaptive_max_pool2d(chunk, size)
|
|
chunk = self.mfr[level](chunk)
|
|
if level:
|
|
chunk = functional.interpolate(
|
|
chunk,
|
|
size=(height, width),
|
|
mode="nearest",
|
|
)
|
|
outputs.append(chunk)
|
|
return self.act(self.aggr(torch.cat(outputs, dim=1))) * value
|
|
|
|
class CCM(nn.Module):
|
|
def __init__(self, channels):
|
|
super().__init__()
|
|
self.ccm = nn.Sequential(
|
|
nn.Conv2d(channels, channels * 2, 3, 1, 1),
|
|
nn.GELU(),
|
|
nn.Conv2d(channels * 2, channels, 1),
|
|
)
|
|
|
|
def forward(self, value):
|
|
return self.ccm(value)
|
|
|
|
class AttBlock(nn.Module):
|
|
def __init__(self, channels):
|
|
super().__init__()
|
|
self.norm1 = LayerNorm(channels)
|
|
self.norm2 = LayerNorm(channels)
|
|
self.safm = SAFM(channels)
|
|
self.ccm = CCM(channels)
|
|
|
|
def forward(self, value):
|
|
value = self.safm(self.norm1(value)) + value
|
|
return self.ccm(self.norm2(value)) + value
|
|
|
|
class SAFMN(nn.Module):
|
|
def __init__(self):
|
|
super().__init__()
|
|
self.to_feat = nn.Conv2d(3, 128, 3, 1, 1)
|
|
self.feats = nn.Sequential(*(AttBlock(128) for _ in range(16)))
|
|
self.to_img = nn.Sequential(
|
|
nn.Conv2d(128, 3 * 4**2, 3, 1, 1),
|
|
nn.PixelShuffle(4),
|
|
)
|
|
|
|
def forward(self, value):
|
|
value = self.to_feat(value)
|
|
return self.to_img(self.feats(value) + value)
|
|
|
|
model = SAFMN()
|
|
model.load_state_dict(
|
|
_load_weights(torch, checkpoint_path, "params_ema"),
|
|
strict=True,
|
|
)
|
|
sample = torch.rand(1, 3, 256, 256)
|
|
_export(torch, model, sample, destination)
|
|
_validate_export(onnx, onnxruntime, torch, model, sample, destination)
|
|
|
|
|
|
def export_model(exporter, checkpoint_path, destination):
|
|
exporters = {
|
|
"realesrgan_wdn_x4": export_realesrgan_wdn_x4,
|
|
"safmn_real_x4": export_safmn_real_x4,
|
|
}
|
|
try:
|
|
export = exporters[exporter]
|
|
except KeyError as error:
|
|
raise ValueError(f"Unknown model exporter: {exporter}") from error
|
|
export(checkpoint_path, destination)
|