| from __future__ import annotations |
|
|
| from pathlib import Path |
| import numpy as np |
| import torch |
|
|
|
|
| class OrtDitModule(torch.nn.Module): |
| """Drop-in replacement for NitroGen's repeated DiT denoise step. |
| |
| CUDA tensors are bound directly to ONNX Runtime buffers, avoiding |
| GPU->CPU->GPU copies. TensorRT EP is used when available; otherwise CUDA EP. |
| """ |
|
|
| def __init__(self, onnx_path: str | Path, *, prefer_tensorrt: bool = True, cache_dir: str | Path = ".ort-cache"): |
| super().__init__() |
| import onnxruntime as ort |
|
|
| try: |
| ort.preload_dlls(directory="") |
| except Exception: |
| try: |
| ort.preload_dlls() |
| except Exception: |
| pass |
|
|
| available = set(ort.get_available_providers()) |
| cache_dir = Path(cache_dir) |
| cache_dir.mkdir(parents=True, exist_ok=True) |
| providers = [] |
| if prefer_tensorrt and "TensorrtExecutionProvider" in available: |
| providers.append(("TensorrtExecutionProvider", { |
| "device_id": 0, |
| "trt_fp16_enable": True, |
| "trt_engine_cache_enable": True, |
| "trt_engine_cache_path": str(cache_dir / "trt"), |
| "trt_timing_cache_enable": True, |
| "trt_timing_cache_path": str(cache_dir / "trt-timing"), |
| "trt_builder_optimization_level": 5, |
| "trt_max_workspace_size": 1073741824, |
| })) |
| if "CUDAExecutionProvider" in available: |
| providers.append(("CUDAExecutionProvider", { |
| "device_id": 0, |
| "arena_extend_strategy": "kSameAsRequested", |
| "do_copy_in_default_stream": True, |
| "cudnn_conv_use_max_workspace": True, |
| "gpu_mem_limit": 1610612736, |
| })) |
| providers.append("CPUExecutionProvider") |
|
|
| opts = ort.SessionOptions() |
| opts.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL |
| opts.enable_mem_pattern = True |
| opts.enable_cpu_mem_arena = True |
| self.session = ort.InferenceSession(str(onnx_path), sess_options=opts, providers=providers) |
| self.provider = self.session.get_providers()[0] |
| self.output_name = self.session.get_outputs()[0].name |
| self._io = self.session.io_binding() |
|
|
| def extra_repr(self) -> str: |
| return f"provider={self.provider}" |
|
|
| def forward(self, hidden_states, encoder_hidden_states, timestep, **_kwargs): |
| h = hidden_states.contiguous().to(dtype=torch.float16) |
| e = encoder_hidden_states.contiguous().to(dtype=torch.float16) |
| t = timestep.contiguous().to(dtype=torch.int64) |
|
|
| if h.is_cuda and self.provider != "CPUExecutionProvider": |
| out = torch.empty(h.shape, device=h.device, dtype=torch.float16) |
| io = self._io |
| io.clear_binding_inputs() |
| io.clear_binding_outputs() |
| io.bind_input("hidden_states", "cuda", h.device.index or 0, np.float16, tuple(h.shape), h.data_ptr()) |
| io.bind_input("encoder_hidden_states", "cuda", e.device.index or 0, np.float16, tuple(e.shape), e.data_ptr()) |
| io.bind_input("timestep", "cuda", t.device.index or 0, np.int64, tuple(t.shape), t.data_ptr()) |
| io.bind_output(self.output_name, "cuda", out.device.index or 0, np.float16, tuple(out.shape), out.data_ptr()) |
| self.session.run_with_iobinding(io) |
| return out |
|
|
| pred = self.session.run([self.output_name], { |
| "hidden_states": h.detach().cpu().numpy(), |
| "encoder_hidden_states": e.detach().cpu().numpy(), |
| "timestep": t.detach().cpu().numpy(), |
| })[0] |
| return torch.from_numpy(pred).to(hidden_states.device, dtype=torch.float16) |
|
|
