"""
ONNX decoder shim for terrain-diffusion.
WHY THE DECODER, AND ONLY THE DECODER
Profiled, not guessed. For one 512^2 tile:
decoder_model 4.39 s 8 calls 549 ms 59.1% of total
base_model 2.67 s 6 calls 445 ms 36.0%
model forward 95.0%
python / tiling / copies 5.0%
95% of the time is inside model forwards, so a faster runtime is aimed at the
real cost rather than at overhead -- and the decoder alone is 59% of it. One
model to shim instead of three, for most of the benefit.
WHY A SHIM AT ALL
The repo ships an ONNX *exporter* and a published ONNX model, but no inference
path: WorldPipeline calls self.decoder_model(x, noise_labels, conditional_inputs)
directly. This replaces that module's forward with an onnxruntime session while
leaving the module object intact, so .to(), .config and everything else the
pipeline touches still work.
ZERO COPY OR IT IS POINTLESS
The naive version moves the tensor GPU -> CPU -> ORT -> CPU -> GPU per call. At
549 ms/call and 8 calls a tile that would swamp any kernel gain. IO binding
passes torch's device pointers straight to ORT, so nothing crosses PCIe.
The published graph has dynamic spatial dims (['batch', 5, 'height', 'width']),
so it works at any decoder_tile_size rather than only the exported one.
"""
import os
import numpy as np
import torch
def _providers(kind):
if kind == "tensorrt":
return [("TensorrtExecutionProvider",
{"trt_fp16_enable": True,
"trt_engine_cache_enable": True,
"trt_engine_cache_path": "/tmp/td_trt_cache"}),
"CUDAExecutionProvider", "CPUExecutionProvider"]
if kind == "cuda":
return ["CUDAExecutionProvider", "CPUExecutionProvider"]
return ["CPUExecutionProvider"]
class OnnxDecoderRunner:
def __init__(self, path, kind="cuda", device=0):
import onnxruntime as ort
self.ort = ort
so = ort.SessionOptions()
so.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
self.sess = ort.InferenceSession(path, so, providers=_providers(kind))
self.device = device
self.kind = kind
got = self.sess.get_providers()
print(f" ONNX decoder: {os.path.basename(path)} providers={got}")
if kind == "tensorrt" and "TensorrtExecutionProvider" not in got:
print(" WARNING: TensorRT provider not active; fell back.")
@torch.no_grad()
def __call__(self, x, noise_labels=None, conditional_inputs=None):
# The graph is fp32. Under autocast x may arrive as fp16, so normalise
# rather than silently feeding a mismatched dtype.
xt = x.detach().to(torch.float32).contiguous()
nt = noise_labels.detach().to(torch.float32).contiguous()
if not xt.is_cuda:
xt = xt.cuda()
if not nt.is_cuda:
nt = nt.cuda()
b, _c, h, w = xt.shape
out = torch.empty((b, 1, h, w), dtype=torch.float32, device=xt.device)
io = self.sess.io_binding()
io.bind_input("x", "cuda", self.device, np.float32, tuple(xt.shape), xt.data_ptr())
io.bind_input("noise_labels", "cuda", self.device, np.float32,
tuple(nt.shape), nt.data_ptr())
io.bind_output("output", "cuda", self.device, np.float32,
tuple(out.shape), out.data_ptr())
self.sess.run_with_iobinding(io)
return out
def install(pipeline, path=None, kind=None):
"""Swap the pipeline's decoder forward for an ONNX session. Returns True if
installed. Controlled by TD_ONNX_DECODER (path) and TD_ONNX_EP."""
path = path or os.environ.get("TD_ONNX_DECODER")
if not path or not os.path.exists(path):
return False
kind = kind or os.environ.get("TD_ONNX_EP", "cuda")
runner = OnnxDecoderRunner(path, kind=kind)
pipeline.decoder_model.forward = runner
pipeline._onnx_decoder = runner # keep a reference alive
return True