---
license: apache-2.0
base_model: nvidia/PiD
tags:
- image-to-image
- super-resolution
- upscaler
- onnx
- tensorrt
- pid
- nvidia
- sdxl
- cpp
pipeline_tag: image-to-image
---

# NVIDIA PiD (Pixel Diffusion) SDXL 4-Step Upscaler — ONNX / TensorRT Version

This repository contains the optimized **ONNX export** of the official NVIDIA PiD (Pixel Diffusion) model: **`PiD_res2kto4k_sr4x_official_sdxl_distill_4step`**. 

It is designed for high-performance super-resolution upscaling (from 1024x1024 to 4096x4096px) in C++ environments using **NVIDIA TensorRT** or **ONNX Runtime**.

## Model Details
* **Original Model:** [NVIDIA PiD (nvidia/PiD)](https://huggingface.co/nvidia/PiD)
* **Upscale Factor:** 4x (e.g., 1024x1024 -> 4096x4096px)
* **Diffusion Steps:** 4 steps (distilled student model)
* **Backbone:** Stable Diffusion XL (SDXL) VAE
* **License:** [Apache-2.0](https://www.apache.org/licenses/LICENSE-2.0) (inherited from the original NVIDIA PiD repository)

---

## What is in this repository?

1. **`vae_encoder.onnx`**
   * The SDXL VAE Encoder wrapper. It converts a preprocessed `1024x1024` input image to `1x4x128x128` latents (`lq_latent`).
   * **Note:** Includes the `quant_conv` projection layer properly to ensure correct latent conditioning without color distortion or noise.

2. **`pid_upscaler.onnx` & `pid_upscaler.onnx.data`**
   * The core PixelDiT student model.
   * **Optimization:** Bypasses the Gemma-2-2B text encoder by using pre-computed zero-prompt embeddings during export, allowing execution in pure C++/TensorRT without Hugging Face gated license checks.
   * **ONNX Compatibility Patches:** Replaced unsupported `torch.nn.functional.fold` operations (Col2Im) with equivalent TensorRT-friendly reshape/permute layers, and fixed float64/float32 type mismatches in sinus-cosinus position embeddings.

---

## TensorRT Compilation Guide

To build high-performance TensorRT engines, use the `trtexec` utility bundled with your TensorRT installation.

> [!WARNING]
> Standard FP16 precision compilation for SDXL VAE causes numerical overflows, resulting in `NaN` latents and blank output images. PiD compilation in FP32 might fail due to extremely high GPU memory allocation requests in the autotuner (over 400 GB). Follow the precision guidelines below to prevent these issues.

### 1. Compile VAE Encoder (FP32)
Compile the VAE in standard FP32 to ensure numerical stability:
```bash
trtexec --onnx=vae_encoder.onnx --saveEngine=vae_encoder.engine --verbose
```

### 2. Compile PiD Upscaler (BFloat16)
Compile the PiD model in **BFloat16** to prevent both `NaN` overflows and GPU out-of-memory errors during build:
```bash
trtexec --onnx=pid_upscaler.onnx --saveEngine=pid_upscaler.engine --bf16 --verbose
```

---

## Inference Pipeline Summary (C++)

In your C++ execution pipeline (e.g., using OpenCV for I/O and CUDA for memory management):

1. **Preprocess:** Resize the input image to `1024x1024`, convert from BGR to RGB, normalize to $[-1, 1]$, and format to CHW.
2. **VAE Inference:** Run `vae_encoder.engine` on the input image to get `1x4x128x128` latents.
3. **Initialize Noise:** Create a starting noise tensor $x_T$ of size `1x3x4096x4096` using normal distribution.
4. **Diffusion Loop:** Loop for 4 steps with scaled timesteps $t \in [999, 866, 634, 342]$:
   * Execute `pid_upscaler.engine` with inputs:
     * `x`: current noisy high-res image `1x3x4096x4096`
     * `t`: scaled timestep value `1`
     * `lq_latent`: VAE latents `1x4x128x128`
     * `degrade_sigma`: noise degradation value `1` (constant `0.0f` for clean images)
   * The model returns `v_pred`.
   * Update the high-res image $x$ using SDE math:
     $$x_{0\_pred} = x_t - t_{cur} \cdot v_{pred}$$
     $$x_{t\_next} = (1.0 - t_{next}) \cdot x_{0\_pred} + t_{next} \cdot \epsilon$$ (where $\epsilon \sim \mathcal{N}(0, I)$)
5. **Postprocess:** Clamp final pixel values to $[-1, 1]$, convert back to $[0, 255]$ (RGB -> BGR), and save the `4096x4096` image.

---

## Citation & Acknowledgements
If you use this model, please cite the original authors:

```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}
}
```
Special thanks to the original NVIDIA team for releasing the PiD model weights.
```