| """Export a Moebius UNet checkpoint to the WebGPU-safe ONNX graph Card Anvil loads. |
| |
| Adapted from simonw/moebius-web python/reexport_unet_no_conv3d.py: the |
| self-attention pos_conv Conv3d is swapped for the equivalent Conv2d (so the |
| graph compiles on every WebGPU backend), then exported at a fixed 512 x 512 |
| (64 x 64 latent) with the same input/output names as the hosted unet.onnx. |
| |
| usage: python export_unet.py <moebius repo> <checkpoint .bin> <out .onnx> |
| """ |
| import importlib.util |
| import os |
| import sys |
|
|
| import numpy as np |
| import torch |
| from torch import nn |
|
|
| repo, weights, out = (os.path.abspath(p) for p in sys.argv[1:4]) |
| sys.path.insert(0, repo) |
|
|
|
|
| class PosConv2d(nn.Module): |
| """Drop-in for the Conv3d pos_conv; keeps the 5D (b, 1|dim_k, v, hh, ww) layout.""" |
|
|
| def __init__(self, c3: nn.Conv3d): |
| super().__init__() |
| dim_k = c3.out_channels |
| r = c3.kernel_size[1] |
| self.conv = nn.Conv2d(1, dim_k, r, padding=r // 2) |
| with torch.no_grad(): |
| self.conv.weight.copy_(c3.weight.squeeze(2)) |
| self.conv.bias.copy_(c3.bias) |
|
|
| def forward(self, V): |
| b, u, v, hh, ww = V.shape |
| x = V.reshape(b * v, 1, hh, ww) |
| x = self.conv(x) |
| return x.reshape(b, v, -1, hh, ww).permute(0, 2, 1, 3, 4) |
|
|
|
|
| class UNetExport(nn.Module): |
| def __init__(self, removal_model): |
| super().__init__() |
| self.m = removal_model |
|
|
| def forward(self, latent_in9, timesteps, input_ids): |
| return self.m(latent_in9, timesteps, input_ids).sample |
|
|
|
|
| def load_removal_module(): |
| |
| |
| path = os.path.join(repo, "removal", "v1_2", "removal_model.py") |
| spec = importlib.util.spec_from_file_location("removal_model", path) |
| module = importlib.util.module_from_spec(spec) |
| spec.loader.exec_module(module) |
| return module |
|
|
|
|
| def main(): |
| rm = load_removal_module() |
| cfg = rm.load_cfg(os.path.join(repo, "config", "model_cfg", "moebius.yaml")) |
| model = rm.build_removal_model(cfg, 20).eval() |
| state = torch.load(weights, map_location="cpu", weights_only=True) |
| print("[load]", model.load_state_dict(state, strict=True)) |
| print(f"[load] params = {sum(p.numel() for p in model.parameters()) / 1e6:.2f}M") |
|
|
| torch.manual_seed(0) |
| latent9 = torch.randn(2, 9, 64, 64) |
| ts = torch.tensor([811, 811], dtype=torch.int64) |
| ids = torch.tensor([list(range(10, 20)), list(range(10))], dtype=torch.int64) |
|
|
| wrap = UNetExport(model) |
| with torch.no_grad(): |
| before = wrap(latent9, ts, ids).numpy() |
| n = 0 |
| for mod in model.modules(): |
| pc = getattr(mod, "pos_conv", None) |
| if isinstance(pc, nn.Conv3d): |
| mod.pos_conv = PosConv2d(pc) |
| n += 1 |
| with torch.no_grad(): |
| after = wrap(latent9, ts, ids).numpy() |
| d = np.abs(before - after).max() |
| print(f"[fix] replaced {n} Conv3d -> Conv2d; torch max|diff| = {d:.3e}") |
| assert n > 0 and d < 1e-4, "Conv2d replacement changed the math" |
|
|
| os.makedirs(os.path.dirname(out), exist_ok=True) |
| torch.onnx.export( |
| wrap, |
| (latent9, ts, ids), |
| out, |
| opset_version=18, |
| dynamo=False, |
| input_names=["latent", "timesteps", "input_ids"], |
| output_names=["noise"], |
| dynamic_axes={ |
| "latent": {0: "B"}, |
| "timesteps": {0: "B"}, |
| "input_ids": {0: "B"}, |
| "noise": {0: "B"}, |
| }, |
| ) |
|
|
| import onnx |
| import onnxruntime as ort |
|
|
| sess = ort.InferenceSession(out, providers=["CPUExecutionProvider"]) |
| got = sess.run( |
| None, |
| {"latent": latent9.numpy(), "timesteps": ts.numpy(), "input_ids": ids.numpy()}, |
| )[0] |
| print(f"[check] ONNX vs torch max|diff| = {np.abs(after - got).max():.3e}") |
| for i in sess.get_inputs() + sess.get_outputs(): |
| print(f"[check] {i.name}: {i.type} {i.shape}") |
| graph = onnx.load(out, load_external_data=False) |
| conv3d = sum( |
| 1 |
| for node in graph.graph.node |
| if node.op_type == "Conv" |
| and any(a.name == "kernel_shape" and len(a.ints) == 3 for a in node.attribute) |
| ) |
| print(f"[check] 3D Conv nodes left: {conv3d}") |
| print(f"[done] {out} ({os.path.getsize(out):,} bytes)") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
