markdown

PiD TensorRT C++ Upscaler 🚀

This repository provides a high-performance C++ inference pipeline for NVIDIA PiD (Pixel Diffusion) using NVIDIA TensorRT and OpenCV. It allows you to upscale images from 1024x1024 to 4096x4096px (4x super-resolution) in a few seconds directly on your GPU. Pre-exported ONNX models are available on Hugging Face: Glebk/PiD-res2kto4k-ONNX.

Features

• Fast C++ Execution: Bypasses Python overhead by implementing the 4-step SDE diffusion sampling loop in pure C++.
• Optimized Model Export:
  • Gemma text encoder is omitted and replaced with pre-computed zero-prompt embeddings to bypass HF license requirements and optimize inference speed.
  • Unsupported PyTorch operations like torch.nn.functional.fold (Col2Im) are patched with TensorRT-friendly layers.
• Robust Precision Management: Prevents NaN overflows in VAE attention layers and compiles the heavy PiD model in BFloat16 (--bf16) to fit inside ordinary GPU memory during build.

---

Project Structure

├── CMakeLists.txt         # CMake configuration file
├── main.cpp               # C++ Inference pipeline (OpenCV + TensorRT)
├── export_vae.py          # Script to export SDXL VAE Encoder with quant_conv to ONNX
├── export_pid.py          # Script to export patched PiD model to ONNX
└── README.md              # Project documentation
Prerequisites
1. Hardware & OS
OS: Windows 10/11 or Linux
GPU: NVIDIA GPU with BFloat16 support (Ampere architecture or newer: RTX 30xx, 40xx, A-series, etc.)
2. Software Dependencies
CUDA Toolkit (>= 12.0)
TensorRT (>= 10.0)
OpenCV (>= 4.x)
CMake (>= 3.18)
C++ Compiler (MSVC on Windows / GCC on Linux) with C++17 support.
Step 1: Exporting Models to ONNX
You can either download the pre-exported ONNX models from Hugging Face or export them yourself.

Patches Required for PiD Export
Before exporting the PiD model, you must apply two fixes to the original nvidia/PiD code:

Fix cos Type Mismatch: In PiD_repo/pid/_src/networks/pixeldit_official.py, inside the get_1d_sincos_pos_embed_from_grid function, cast the inputs of torch.cos and torch.sin to float32 to prevent export failures.

Replace torch.nn.functional.fold: TensorRT does not support fold with dynamic spatial dimensions. In PiD_repo/pid/_src/networks/pid_net.py (or corresponding module), replace the F.fold operation on patches back to the image space with an equivalent sequence of reshape and permute operations.

Run Export Scripts:
bash


# 1. Export VAE Encoder
python export_vae.py --model stabilityai/sdxl-vae --output vae_encoder.onnx
# 2. Export PiD Student Model
python export_pid.py --checkpoint /path/to/model_ema_bf16.pth --output pid_upscaler.onnx
Step 2: Compiling Engines via trtexec
Compiling these models with correct precision settings is critical to prevent numerical overflows (NaN outputs) and compilation Out-Of-Memory errors.

1. Compile VAE (FP32)
Do not use --fp16 for VAE, otherwise, its attention blocks will overflow and return NaNs.

bash


trtexec --onnx=vae_encoder.onnx --saveEngine=vae_encoder.engine --verbose
2. Compile PiD Upscaler (BFloat16)
Using --bf16 ensures FP32-like range to prevent NaNs while keeping memory consumption low enough to compile successfully on standard GPUs.

bash


trtexec --onnx=pid_upscaler.onnx --saveEngine=pid_upscaler.engine --bf16 --verbose
Step 3: Building the C++ Application
Configure environment variables for CMake to find your dependencies (e.g., set TENSORRT_PATH and OpenCV_DIR).
Generate build files and compile:
bash


mkdir build
cd build
cmake .. -DTENSORRT_PATH="C:/path/to/TensorRT"
cmake --build . --config Release
Step 4: Running the Upscaler
Run the compiled executable passing the paths to your test image and the generated TensorRT engines:

bash


./Release/pid_upscaler.exe <input_image> vae_encoder.engine pid_upscaler.engine <output_image> [seed]
Example:
bash


./Release/pid_upscaler.exe FOTO1532.jpg vae_encoder.engine pid_upscaler.engine output_4k.png 42
Citation & References
Original PiD Repository: nvidia/PiD

bibtex


@article{lu2025pid,
  title={PiD: Halving the Steps of Pixel-space Diffusion Upscalers via Distillation},
  author={Lu, Yifan and others},
  journal={arXiv preprint arXiv:2501.00000},
  year={2025}
}