CorridorKey ONNX Graphs

These are weight-free ONNX graphs of the real-time CorridorKey 2.0 pipeline (MoGe-2 encoder, RVM step, Stage 1, Stage 2 and the CorridorKey Edge Refiner), made for the CorridorKey OBS plugin. They contain the network structure and small constants only. Every trained weight is downloaded from its original publisher and bound locally, using the manifest that comes with each graph.

_{Real outputs on one frame of a Pexels clip (Pexels licence), 256x512 grid. The right column is the absolute
difference multiplied by 20.}

Files

graphs/<name>.onnx is the graph and graphs/<name>.json its manifest. Two inference grids are included: 256x512 (the facecam profile, full frame) and 512x512 (person crop).

Graph	Inputs	Outputs	Weights from
`moge_256x512`, `moge_512x512`	`rgb_moge` (1,3,H·7/32,W·7/32)	`moge_feat` (1,1,768,H/64,W/64)	MoGe-2 encoder, CK stats
`rvm_256x512`, `rvm_512x512`	`rgb_rvm` (1,3,H/4,W/4), `state` (1,1+H/64·W/64,384)	`rvm_feat` (1,1,384,H/64,W/64), `state_out`	RVM small, CK stats
`stage1nomask_256x512`, `stage1nomask_512x512`	`rgb` (1,3,H,W), `moge_feat`, `rvm_feat`, `hint` (1,1,1,H,W), `hint_quality` (1,1)	`baton` (1,1,512,H/16,W/16), `alpha` (1,1,H/2,W/2), `fg` (1,3,H/2,W/2)	CorridorKey 2.0
`stage2nomask_256x512`, `stage2nomask_512x512`	`rgb` (1,3,H,W), `baton`	`alpha` (1,1,H,W), `fg` (1,3,H,W)	CorridorKey 2.0
`stage1rvmnomask_256x512`, `stage1rvmnomask_512x512` (optional)	`rgb`, `state`, `moge_feat`, `hint`, `hint_quality`	`baton`, `alpha`, `fg`, `state_out`	CorridorKey 2.0, RVM, CK stats
`refiner_b256`, `_b384`, `_b512`, `_b768`	`x` (N,11,52,52)	`alpha` (N,1,52,52), `fg` (N,3,52,52)	Edge Refiner v12

All inputs and outputs are float32. rgb, rgb_moge and rgb_rvm are sRGB in [0, 1]. The MoGe and RVM inputs are box-filter (area) downsamples of the rgb grid image. fg is straight (not premultiplied) sRGB in [0, 1]; the sRGB to linear to asinh conversion and its inverse happen inside the graphs. With no hint, hint is zeros and hint_quality is -1. state starts as zeros and is carried from state_out to the next frame.

The manifest

Each .onnx refers to its weights as standard ONNX external data in <name>.weights, which is not in this repository. The manifest describes how to build that file:

inputs and outputs: name, dtype, shape.
weights: one entry per initializer with its name, dtype, shape, offset and length in <name>.weights (64-byte aligned), and source = {file, key, ops}.
sources: the pinned downloads (repository or URL, revision, file name, size, sha256, licence).

ops turn the stored array into the initializer, in order: ["squeeze", axis], ["transpose", perm], ["merge", a, b] (merge axes [a, b) into one, C order) and ["zeros"] (no source tensor). Only the RVM weights need ops, because the original checkpoint is in JAX layout. The C++ binder of the OBS plugin and the Python reference (tools/onnx_weights.py) produce byte-identical .weights files for all graphs.

Weight sources

id	File	From	Licence
`ck2`, `stats`	`corridorkey_v2.safetensors`, `foundation_stats.safetensors`	nikopueringer/CorridorKey2.0	Corridor Key Licence
`moge`	`model.safetensors`	mantelmoritz/MoGe-2-ViT-B-Encoder-safetensors (format conversion of Ruicheng/moge-2-vitb-normal)	MIT (Microsoft), Apache-2.0 (DINOv2)
`rvm`	`pretrain_rvm_small16_256_204031069.npz`	Google DeepMind representations4d	CC BY 4.0
`refiner`	`corridorkey_edge_refiner_v12.safetensors`	mantelmoritz/CorridorKey-Realtime-EdgeRefiner	Corridor Key Licence

The exact revisions and sha256 sums are in each manifest.

How the graphs differ from upstream

The graphs were exported from the upstream PyTorch modules with torch.onnx.export(dynamo=True) at opset 18, in float32. These changes are exact, apart from floating-point summation order:

stage1nomask and stage2nomask run the shifted Swin windows without region masks, as the real-time engine does. Upstream applies the masks. The Edge Refiner was trained on the mask-free Stage-1 outputs, and with masks the real-time pipeline measured 3 to 6 % higher edge error against the offline teacher.
The 7x7 depthwise SwiGLU convolutions run as two depthwise convolutions (value and gate) instead of one grouped convolution with two filters per channel. DirectML rejects the grouped form.
The depatchify ConvTranspose (kernel 2s, stride s) is a 2x2 convolution, DepthToSpace and a crop. On TensorRT fp16 this makes Stage 2 about 15 % faster at the same accuracy. The idea comes from the export notes of corridorkey-desktop; the code here is a separate implementation.
Graph cleanup with the onnxscript optimizer and ONNX shape inference. The weights were graph inputs during this step, so no weight was folded into a constant.
For DirectML: Reshape without allowzero, two-sided Clip as Min/Max, and rank-5 MatMul inputs flattened to rank 2 and reshaped back.
alpha and fg are decoded inside the graph, so a consumer needs no CorridorKey colour constants.

Measured

RTX 5070 Ti, ONNX Runtime 1.31.1 with TensorRT 10 (TensorRT execution provider) and ONNX Runtime DirectML 1.24.4. Device time per call with inputs and outputs bound on the GPU, median of 40 runs (DirectML: 10). Errors are against the PyTorch fp32 modules on a real facecam frame. The desktop was in use during the runs, so expect about ±1 ms of noise.

Graph (256x512)	TensorRT fp16	TensorRT fp32	DirectML fp32	alpha error, TensorRT fp16 (max / mean)
`stage1nomask`	7 to 8 ms	10.4 ms	15.6 ms	8.0e-3 / 3.2e-4
`stage2nomask`	3.8 ms	8.6 ms	18.1 ms	1.0e-2 / 2.4e-4
`rvm`	2.1 ms	2.8 ms	3.4 ms	(feature error 8.5e-3 of range in fp16)
`moge`	2.0 ms	2.6 ms	3.3 ms	(feature error 4.5e-2 of range in fp16)
`refiner_b384`	2.0 ms	4.6 ms	15.2 ms	2.1e-3 / 2.5e-4

Recommended precisions: fp16 for Stage 1, Stage 2 and the refiner; fp32 for RVM and MoGe. In fp16 the MoGe features drift by up to 4.5 % of their range, and RVM's error builds up through the recurrent state. bf16 was worse than fp16 in every graph and does not build for Stage 1 and Stage 2. The CUDA execution provider without TensorRT works (fp32 error about 1e-3) but is 4 to 8 times slower.

At 512x512: stage1nomask 8.0 ms, stage2nomask 6.5 ms on TensorRT fp16; 21.8 ms and 34.5 ms on DirectML.

With CUDA graph capture

The TensorRT execution provider can capture each engine as a CUDA graph (trt_cuda_graph_enable=1). Inputs and outputs then have to stay bound to the same GPU buffers. Most of the time above is kernel launch overhead, so capture helps a lot, with identical outputs:

Graph	Precision	Without capture	With capture
`rvm_256x512`	fp32	2.24 ms	0.87 ms
`rvm_512x512`	fp32	3.38 ms	0.89 ms
`stage1nomask_256x512`	fp16	7.29 ms	3.81 ms
`stage1nomask_512x512`	fp16	10.42 ms	4.73 ms
`stage2nomask_256x512`	fp16	4.13 ms	3.39 ms
`refiner_b384`	fp16	1.87 ms	1.80 ms

moge_* does not run under capture yet; it is only needed every few frames.

Changes

Revision 3: moge_* without the antialiased Resize to the size the graph is already fed at (DirectML ran that node on the CPU). rvm_* stacks the per-step output instead of writing it into an empty tensor, which exported as ScatterElements and also ran on the CPU under DirectML. Both graphs now run entirely on DirectML, with unchanged outputs and weights. New optional stage1rvmnomask_*: the RVM step and Stage 1 in one graph, with the RVM input computed in the graph as the 4x4 average of rgb. It saves one call per frame, but a weakly typed TensorRT fp16 build also runs RVM in fp16, so the separate graphs stay the default.
Revision 2: stage1nomask and stage2nomask carry their RoPE tables in float32. Revision 1 had bfloat16 tables by mistake. Parity against the PyTorch fp32 reference stays in the same range (TensorRT fp32 alpha max 8e-4, fp16 8e-3). The manifests and the weights are unchanged.
Revision 1: first release.

Known limits

The graphs need a fixed grid. Other grid sizes need a new export.
DirectML runs everything, but Stage 1 and Stage 2 are 2 to 5 times slower than TensorRT on the same card. AMD and MIGraphX were not tested.
The parity numbers come from one frame per grid. The end-to-end comparison on the evaluation clips runs in the PyTorch engine of the research repository.

Licence and attribution

The graphs are derived from CorridorKey 2.0 and published under the Corridor Key Licence, Version 2.0 (CC BY-NC-SA 4.0 with Additional Terms by Corridor Digital LLC). The full text is in LICENSE. Derivatives must stay free of charge, use the same licence and keep the CorridorKey name. Commercial software integrations need a separate agreement with Corridor Digital (contact@corridordigital.com).

CorridorKey 2.0 © Corridor Digital LLC, designed and trained by Zachary Harris in collaboration with Niko Pueringer and Corridor Digital: https://huggingface.co/nikopueringer/CorridorKey2.0, https://corridorkey.com
MoGe-2 (Microsoft, MIT) with a DINOv2 backbone (Meta, Apache-2.0): https://github.com/microsoft/MoGe
RVM (Google DeepMind; code Apache-2.0, weights CC BY 4.0): https://github.com/google-deepmind/representations4d
No weights of any of these models are contained in this repository.
This is an independent community project. It is not affiliated with, sponsored by or endorsed by Corridor Digital LLC, Microsoft, Meta or Google DeepMind.

Downloads last month: -; Downloads are not tracked for this model. How to track