github-featured/nvidia--tensorrt
2
0
Fork 0
Code Issues Pull Requests Deployments Models Agent Notes Packages Projects Releases Wiki Activity

TensorRT 10.9 OSS Release. (#4381)

Browse Source 
Signed-off-by: Leo Dong <leod@nvidia.com>
This commit is contained in:
Leo Dong
2025-03-11 14:54:52 -07:00
committed by GitHub
parent 64e56ab6cb
commit 8c6d69ddec
474 changed files with 12220 additions and 6843 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.gitignore
-1
View File
@@ -3,7 +3,6 @@ build/
/demo/BERT/engines
/demo/BERT/squad/*.json
/docker/jetpack_files/*
*.nvmk
*.sln
*.vcxproj
externals/
CHANGELOG.md
+318 -185
View File
File diff suppressed because it is too large Load Diff
CMakeLists.txt
+11 -25
View File
@@ -1,5 +1,5 @@
#
# SPDX-FileCopyrightText: Copyright (c) 1993-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-FileCopyrightText: Copyright (c) 1993-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#
# Licensed under the Apache License, Version 2.0 (the "License");
@@ -176,35 +176,21 @@ set(CUDA_LIBRARIES ${CUDART_LIB})
if (DEFINED GPU_ARCHS)
message(STATUS "GPU_ARCHS defined as ${GPU_ARCHS}. Generating CUDA code for SM ${GPU_ARCHS}")
separate_arguments(GPU_ARCHS)
foreach(SM IN LISTS GPU_ARCHS)
list(APPEND CMAKE_CUDA_ARCHITECTURES SM)
endforeach()
else()
list(APPEND GPU_ARCHS
75
)
find_file(IS_L4T_NATIVE nv_tegra_release PATHS /env/)
set (IS_L4T_CROSS "False")
if (DEFINED ENV{IS_L4T_CROSS})
set(IS_L4T_CROSS $ENV{IS_L4T_CROSS})
list(APPEND CMAKE_CUDA_ARCHITECTURES 72 75 80 86 87 89 90)
if(CUDA_VERSION VERSION_GREATER_EQUAL 12.8)
list(APPEND CMAKE_CUDA_ARCHITECTURES 100 120)
endif()
if (IS_L4T_NATIVE OR ${IS_L4T_CROSS} STREQUAL "True")
# Only Orin (SM87) supported
list(APPEND GPU_ARCHS 87)
endif()
if (CUDA_VERSION VERSION_GREATER_EQUAL 11.0)
# Ampere GPU (SM80) support is only available in CUDA versions > 11.0
list(APPEND GPU_ARCHS 80)
endif()
if (CUDA_VERSION VERSION_GREATER_EQUAL 11.1)
list(APPEND GPU_ARCHS 86)
endif()
message(STATUS "GPU_ARCHS is not defined. Generating CUDA code for default SMs: ${GPU_ARCHS}")
message(STATUS "GPU_ARCHS is not defined. Generating CUDA code for default SMs: ${CMAKE_CUDA_ARCHITECTURES}")
endif()
set(BERT_GENCODES)
# Generate SASS for each architecture
foreach(arch ${GPU_ARCHS})
foreach(arch ${CMAKE_CUDA_ARCHITECTURES})
if (${arch} GREATER_EQUAL 75)
set(BERT_GENCODES "${BERT_GENCODES} -gencode arch=compute_${arch},code=sm_${arch}")
endif()
@@ -212,7 +198,7 @@ foreach(arch ${GPU_ARCHS})
endforeach()
# Generate PTX for the last architecture in the list.
list(GET GPU_ARCHS -1 LATEST_SM)
list(GET CMAKE_CUDA_ARCHITECTURES -1 LATEST_SM)
set(GENCODES "${GENCODES} -gencode arch=compute_${LATEST_SM},code=compute_${LATEST_SM}")
if (${LATEST_SM} GREATER_EQUAL 75)
set(BERT_GENCODES "${BERT_GENCODES} -gencode arch=compute_${LATEST_SM},code=compute_${LATEST_SM}")
README.md
+164 -143
View File
@@ -1,210 +1,231 @@
[![License](https://img.shields.io/badge/License-Apache%202.0-blue.svg)](https://opensource.org/licenses/Apache-2.0) [![Documentation](https://img.shields.io/badge/TensorRT-documentation-brightgreen.svg)](https://docs.nvidia.com/deeplearning/sdk/tensorrt-developer-guide/index.html) [![Roadmap](https://img.shields.io/badge/Roadmap-Q1_2025-brightgreen.svg)](documents/tensorrt_roadmap_2025q1.pdf)
# TensorRT Open Source Software
This repository contains the Open Source Software (OSS) components of NVIDIA TensorRT. It includes the sources for TensorRT plugins and ONNX parser, as well as sample applications demonstrating usage and capabilities of the TensorRT platform. These open source software components are a subset of the TensorRT General Availability (GA) release with some extensions and bug-fixes.
* For code contributions to TensorRT-OSS, please see our [Contribution Guide](CONTRIBUTING.md) and [Coding Guidelines](CODING-GUIDELINES.md).
* For a summary of new additions and updates shipped with TensorRT-OSS releases, please refer to the [Changelog](CHANGELOG.md).
* For business inquiries, please contact [researchinquiries@nvidia.com](mailto:researchinquiries@nvidia.com)
* For press and other inquiries, please contact Hector Marinez at [hmarinez@nvidia.com](mailto:hmarinez@nvidia.com)
- For code contributions to TensorRT-OSS, please see our [Contribution Guide](CONTRIBUTING.md) and [Coding Guidelines](CODING-GUIDELINES.md).
- For a summary of new additions and updates shipped with TensorRT-OSS releases, please refer to the [Changelog](CHANGELOG.md).
- For business inquiries, please contact [researchinquiries@nvidia.com](mailto:researchinquiries@nvidia.com)
- For press and other inquiries, please contact Hector Marinez at [hmarinez@nvidia.com](mailto:hmarinez@nvidia.com)
Need enterprise support? NVIDIA global support is available for TensorRT with the [NVIDIA AI Enterprise software suite](https://www.nvidia.com/en-us/data-center/products/ai-enterprise/). Check out [NVIDIA LaunchPad](https://www.nvidia.com/en-us/launchpad/ai/ai-enterprise/) for free access to a set of hands-on labs with TensorRT hosted on NVIDIA infrastructure.
Join the [TensorRT and Triton community](https://www.nvidia.com/en-us/deep-learning-ai/triton-tensorrt-newsletter/) and stay current on the latest product updates, bug fixes, content, best practices, and more.
# Prebuilt TensorRT Python Package
We provide the TensorRT Python package for an easy installation. \
To install:
```bash
pip install tensorrt
```
You can skip the **Build** section to enjoy TensorRT with Python.
# Build
## Prerequisites
To build the TensorRT-OSS components, you will first need the following software packages.
**TensorRT GA build**
* TensorRT v10.8.0.43
* Available from direct download links listed below
- TensorRT v10.9.0.34
- Available from direct download links listed below
**System Packages**
* [CUDA](https://developer.nvidia.com/cuda-toolkit)
* Recommended versions:
* cuda-12.8.0 + cuDNN-8.9
* cuda-11.8.0 + cuDNN-8.9
* [GNU make](https://ftp.gnu.org/gnu/make/) >= v4.1
* [cmake](https://github.com/Kitware/CMake/releases) >= v3.13
* [python](<https://www.python.org/downloads/>) >= v3.8, <= v3.10.x
* [pip](https://pypi.org/project/pip/#history) >= v19.0
* Essential utilities
* [git](https://git-scm.com/downloads), [pkg-config](https://www.freedesktop.org/wiki/Software/pkg-config/), [wget](https://www.gnu.org/software/wget/faq.html#download)
- [CUDA](https://developer.nvidia.com/cuda-toolkit)
- Recommended versions:
- cuda-12.8.0 + cuDNN-8.9
- cuda-11.8.0 + cuDNN-8.9
- [GNU make](https://ftp.gnu.org/gnu/make/) >= v4.1
- [cmake](https://github.com/Kitware/CMake/releases) >= v3.13
- [python](https://www.python.org/downloads/) >= v3.8, <= v3.10.x
- [pip](https://pypi.org/project/pip/#history) >= v19.0
- Essential utilities
- [git](https://git-scm.com/downloads), [pkg-config](https://www.freedesktop.org/wiki/Software/pkg-config/), [wget](https://www.gnu.org/software/wget/faq.html#download)
**Optional Packages**
* Containerized build
* [Docker](https://docs.docker.com/install/) >= 19.03
* [NVIDIA Container Toolkit](https://github.com/NVIDIA/nvidia-docker)
* PyPI packages (for demo applications/tests)
* [onnx](https://pypi.org/project/onnx/)
* [onnxruntime](https://pypi.org/project/onnxruntime/)
* [tensorflow-gpu](https://pypi.org/project/tensorflow/) >= 2.5.1
* [Pillow](https://pypi.org/project/Pillow/) >= 9.0.1
* [pycuda](https://pypi.org/project/pycuda/) < 2021.1
* [numpy](https://pypi.org/project/numpy/)
* [pytest](https://pypi.org/project/pytest/)
* Code formatting tools (for contributors)
* [Clang-format](https://clang.llvm.org/docs/ClangFormat.html)
* [Git-clang-format](https://github.com/llvm-mirror/clang/blob/master/tools/clang-format/git-clang-format)
- Containerized build
- [Docker](https://docs.docker.com/install/) >= 19.03
- [NVIDIA Container Toolkit](https://github.com/NVIDIA/nvidia-docker)
- PyPI packages (for demo applications/tests)
- [onnx](https://pypi.org/project/onnx/)
- [onnxruntime](https://pypi.org/project/onnxruntime/)
- [tensorflow-gpu](https://pypi.org/project/tensorflow/) >= 2.5.1
- [Pillow](https://pypi.org/project/Pillow/) >= 9.0.1
- [pycuda](https://pypi.org/project/pycuda/) < 2021.1
- [numpy](https://pypi.org/project/numpy/)
- [pytest](https://pypi.org/project/pytest/)
- Code formatting tools (for contributors)
- [Clang-format](https://clang.llvm.org/docs/ClangFormat.html)
- [Git-clang-format](https://github.com/llvm-mirror/clang/blob/master/tools/clang-format/git-clang-format)
> NOTE: [onnx-tensorrt](https://github.com/onnx/onnx-tensorrt), [cub](http://nvlabs.github.io/cub/), and [protobuf](https://github.com/protocolbuffers/protobuf.git) packages are downloaded along with TensorRT OSS, and not required to be installed.
## Downloading TensorRT Build
1. #### Download TensorRT OSS
```bash
git clone -b main https://github.com/nvidia/TensorRT TensorRT
cd TensorRT
git submodule update --init --recursive
```
```bash
git clone -b main https://github.com/nvidia/TensorRT TensorRT
cd TensorRT
git submodule update --init --recursive
```
2. #### (Optional - if not using TensorRT container) Specify the TensorRT GA release build path
If using the TensorRT OSS build container, TensorRT libraries are preinstalled under `/usr/lib/x86_64-linux-gnu` and you may skip this step.
If using the TensorRT OSS build container, TensorRT libraries are preinstalled under `/usr/lib/x86_64-linux-gnu` and you may skip this step.
Else download and extract the TensorRT GA build from [NVIDIA Developer Zone](https://developer.nvidia.com) with the direct links below:
- [TensorRT 10.8.0.43 for CUDA 11.8, Linux x86_64](https://developer.nvidia.com/downloads/compute/machine-learning/tensorrt/10.8.0/tars/TensorRT-10.8.0.43.Linux.x86_64-gnu.cuda-11.8.tar.gz)
- [TensorRT 10.8.0.43 for CUDA 12.8, Linux x86_64](https://developer.nvidia.com/downloads/compute/machine-learning/tensorrt/10.8.0/tars/TensorRT-10.8.0.43.Linux.x86_64-gnu.cuda-12.8.tar.gz)
- [TensorRT 10.8.0.43 for CUDA 11.8, Windows x86_64](https://developer.nvidia.com/downloads/compute/machine-learning/tensorrt/10.8.0/zip/TensorRT-10.8.0.43.Windows.win10.cuda-11.8.zip)
- [TensorRT 10.8.0.43 for CUDA 12.8, Windows x86_64](https://developer.nvidia.com/downloads/compute/machine-learning/tensorrt/10.8.0/zip/TensorRT-10.8.0.43.Windows.win10.cuda-12.8.zip)
Else download and extract the TensorRT GA build from [NVIDIA Developer Zone](https://developer.nvidia.com) with the direct links below:
- [TensorRT 10.9.0.34 for CUDA 11.8, Linux x86_64](https://developer.nvidia.com/downloads/compute/machine-learning/tensorrt/10.9.0/tars/TensorRT-10.9.0.34.Linux.x86_64-gnu.cuda-11.8.tar.gz)
- [TensorRT 10.9.0.34 for CUDA 12.8, Linux x86_64](https://developer.nvidia.com/downloads/compute/machine-learning/tensorrt/10.9.0/tars/TensorRT-10.9.0.34.Linux.x86_64-gnu.cuda-12.8.tar.gz)
- [TensorRT 10.9.0.34 for CUDA 11.8, Windows x86_64](https://developer.nvidia.com/downloads/compute/machine-learning/tensorrt/10.9.0/zip/TensorRT-10.9.0.34.Windows.win10.cuda-11.8.zip)
- [TensorRT 10.9.0.34 for CUDA 12.8, Windows x86_64](https://developer.nvidia.com/downloads/compute/machine-learning/tensorrt/10.9.0/zip/TensorRT-10.9.0.34.Windows.win10.cuda-12.8.zip)
**Example: Ubuntu 20.04 on x86-64 with cuda-12.8**
**Example: Ubuntu 20.04 on x86-64 with cuda-12.8**
```bash
cd ~/Downloads
tar -xvzf TensorRT-10.8.0.43.Linux.x86_64-gnu.cuda-12.8.tar.gz
export TRT_LIBPATH=`pwd`/TensorRT-10.8.0.43
```
```bash
cd ~/Downloads
tar -xvzf TensorRT-10.9.0.34.Linux.x86_64-gnu.cuda-12.8.tar.gz
export TRT_LIBPATH=`pwd`/TensorRT-10.9.0.34
```
**Example: Windows on x86-64 with cuda-12.8**
**Example: Windows on x86-64 with cuda-12.8**
```powershell
Expand-Archive -Path TensorRT-10.8.0.43.Windows.win10.cuda-12.8.zip
$env:TRT_LIBPATH="$pwd\TensorRT-10.8.0.43\lib"
```
```powershell
Expand-Archive -Path TensorRT-10.9.0.34.Windows.win10.cuda-12.8.zip
$env:TRT_LIBPATH="$pwd\TensorRT-10.9.0.34\lib"
```
## Setting Up The Build Environment
For Linux platforms, we recommend that you generate a docker container for building TensorRT OSS as described below. For native builds, please install the [prerequisite](#prerequisites) *System Packages*.
For Linux platforms, we recommend that you generate a docker container for building TensorRT OSS as described below. For native builds, please install the [prerequisite](#prerequisites) _System Packages_.
1. #### Generate the TensorRT-OSS build container.
The TensorRT-OSS build container can be generated using the supplied Dockerfiles and build scripts. The build containers are configured for building TensorRT OSS out-of-the-box.
**Example: Ubuntu 20.04 on x86-64 with cuda-12.8 (default)**
```bash
./docker/build.sh --file docker/ubuntu-20.04.Dockerfile --tag tensorrt-ubuntu20.04-cuda12.8
```
**Example: Rockylinux8 on x86-64 with cuda-12.8**
```bash
./docker/build.sh --file docker/rockylinux8.Dockerfile --tag tensorrt-rockylinux8-cuda12.8
```
**Example: Ubuntu 22.04 cross-compile for Jetson (aarch64) with cuda-12.8 (JetPack SDK)**
```bash
./docker/build.sh --file docker/ubuntu-cross-aarch64.Dockerfile --tag tensorrt-jetpack-cuda12.8
```
**Example: Ubuntu 22.04 on aarch64 with cuda-12.8**
```bash
./docker/build.sh --file docker/ubuntu-22.04-aarch64.Dockerfile --tag tensorrt-aarch64-ubuntu22.04-cuda12.8
```
**Example: Ubuntu 20.04 on x86-64 with cuda-12.8 (default)**
```bash
./docker/build.sh --file docker/ubuntu-20.04.Dockerfile --tag tensorrt-ubuntu20.04-cuda12.8
```
**Example: Rockylinux8 on x86-64 with cuda-12.8**
```bash
./docker/build.sh --file docker/rockylinux8.Dockerfile --tag tensorrt-rockylinux8-cuda12.8
```
**Example: Ubuntu 22.04 cross-compile for Jetson (aarch64) with cuda-12.8 (JetPack SDK)**
```bash
./docker/build.sh --file docker/ubuntu-cross-aarch64.Dockerfile --tag tensorrt-jetpack-cuda12.8
```
**Example: Ubuntu 22.04 on aarch64 with cuda-12.8**
```bash
./docker/build.sh --file docker/ubuntu-22.04-aarch64.Dockerfile --tag tensorrt-aarch64-ubuntu22.04-cuda12.8
```
2. #### Launch the TensorRT-OSS build container.
**Example: Ubuntu 20.04 build container**
```bash
./docker/launch.sh --tag tensorrt-ubuntu20.04-cuda12.8 --gpus all
```
> NOTE:
<br> 1. Use the `--tag` corresponding to build container generated in Step 1.
<br> 2. [NVIDIA Container Toolkit](#prerequisites) is required for GPU access (running TensorRT applications) inside the build container.
<br> 3. `sudo` password for Ubuntu build containers is 'nvidia'.
<br> 4. Specify port number using `--jupyter <port>` for launching Jupyter notebooks.
**Example: Ubuntu 20.04 build container**
```bash
./docker/launch.sh --tag tensorrt-ubuntu20.04-cuda12.8 --gpus all
```
> NOTE:
> <br> 1. Use the `--tag` corresponding to build container generated in Step 1.
> <br> 2. [NVIDIA Container Toolkit](#prerequisites) is required for GPU access (running TensorRT applications) inside the build container.
> <br> 3. `sudo` password for Ubuntu build containers is 'nvidia'.
> <br> 4. Specify port number using `--jupyter <port>` for launching Jupyter notebooks.
## Building TensorRT-OSS
* Generate Makefiles and build.
**Example: Linux (x86-64) build with default cuda-12.8**
```bash
cd $TRT_OSSPATH
mkdir -p build && cd build
cmake .. -DTRT_LIB_DIR=$TRT_LIBPATH -DTRT_OUT_DIR=`pwd`/out
make -j$(nproc)
```
**Example: Linux (aarch64) build with default cuda-12.8**
```bash
cd $TRT_OSSPATH
mkdir -p build && cd build
cmake .. -DTRT_LIB_DIR=$TRT_LIBPATH -DTRT_OUT_DIR=`pwd`/out -DCMAKE_TOOLCHAIN_FILE=$TRT_OSSPATH/cmake/toolchains/cmake_aarch64-native.toolchain
make -j$(nproc)
```
**Example: Native build on Jetson (aarch64) with cuda-12.8**
```bash
cd $TRT_OSSPATH
mkdir -p build && cd build
cmake .. -DTRT_LIB_DIR=$TRT_LIBPATH -DTRT_OUT_DIR=`pwd`/out -DTRT_PLATFORM_ID=aarch64 -DCUDA_VERSION=12.8
- Generate Makefiles and build
**Example: Linux (x86-64) build with default cuda-12.8**
```bash
cd $TRT_OSSPATH
mkdir -p build && cd build
cmake .. -DTRT_LIB_DIR=$TRT_LIBPATH -DTRT_OUT_DIR=`pwd`/out
make -j$(nproc)
```
**Example: Linux (aarch64) build with default cuda-12.8**
```bash
cd $TRT_OSSPATH
mkdir -p build && cd build
cmake .. -DTRT_LIB_DIR=$TRT_LIBPATH -DTRT_OUT_DIR=`pwd`/out -DCMAKE_TOOLCHAIN_FILE=$TRT_OSSPATH/cmake/toolchains/cmake_aarch64-native.toolchain
make -j$(nproc)
```
**Example: Native build on Jetson (aarch64) with cuda-12.8**
```bash
cd $TRT_OSSPATH
mkdir -p build && cd build
cmake .. -DTRT_LIB_DIR=$TRT_LIBPATH -DTRT_OUT_DIR=`pwd`/out -DTRT_PLATFORM_ID=aarch64 -DCUDA_VERSION=12.8
CC=/usr/bin/gcc make -j$(nproc)
```
```
> NOTE: C compiler must be explicitly specified via CC= for native aarch64 builds of protobuf.
**Example: Ubuntu 22.04 Cross-Compile for Jetson (aarch64) with cuda-12.8 (JetPack)**
```bash
cd $TRT_OSSPATH
mkdir -p build && cd build
cmake .. -DCMAKE_TOOLCHAIN_FILE=$TRT_OSSPATH/cmake/toolchains/cmake_aarch64.toolchain -DCUDA_VERSION=12.8 -DCUDNN_LIB=/pdk_files/cudnn/usr/lib/aarch64-linux-gnu/libcudnn.so -DCUBLAS_LIB=/usr/local/cuda-12.8/targets/aarch64-linux/lib/stubs/libcublas.so -DCUBLASLT_LIB=/usr/local/cuda-12.8/targets/aarch64-linux/lib/stubs/libcublasLt.so -DTRT_LIB_DIR=/pdk_files/tensorrt/lib
make -j$(nproc)
```
**Example: Ubuntu 22.04 Cross-Compile for Jetson (aarch64) with cuda-12.8 (JetPack)**
**Example: Native builds on Windows (x86) with cuda-12.8**
```powershell
cd $TRT_OSSPATH
mkdir -p build
cd -p build
cmake .. -DTRT_LIB_DIR="$env:TRT_LIBPATH" -DCUDNN_ROOT_DIR="$env:CUDNN_PATH" -DTRT_OUT_DIR="$pwd\\out"
msbuild TensorRT.sln /property:Configuration=Release -m:$env:NUMBER_OF_PROCESSORS
```
```bash
cd $TRT_OSSPATH
mkdir -p build && cd build
cmake .. -DCMAKE_TOOLCHAIN_FILE=$TRT_OSSPATH/cmake/toolchains/cmake_aarch64.toolchain -DCUDA_VERSION=12.8 -DCUDNN_LIB=/pdk_files/cudnn/usr/lib/aarch64-linux-gnu/libcudnn.so -DCUBLAS_LIB=/usr/local/cuda-12.8/targets/aarch64-linux/lib/stubs/libcublas.so -DCUBLASLT_LIB=/usr/local/cuda-12.8/targets/aarch64-linux/lib/stubs/libcublasLt.so -DTRT_LIB_DIR=/pdk_files/tensorrt/lib
make -j$(nproc)
```
> NOTE:
<br> 1. The default CUDA version used by CMake is 12.4.0. To override this, for example to 11.8, append `-DCUDA_VERSION=11.8` to the cmake command.
* Required CMake build arguments are:
- `TRT_LIB_DIR`: Path to the TensorRT installation directory containing libraries.
- `TRT_OUT_DIR`: Output directory where generated build artifacts will be copied.
* Optional CMake build arguments:
- `CMAKE_BUILD_TYPE`: Specify if binaries generated are for release or debug (contain debug symbols). Values consists of [`Release`] | `Debug`
- `CUDA_VERSION`: The version of CUDA to target, for example [`11.7.1`].
- `CUDNN_VERSION`: The version of cuDNN to target, for example [`8.6`].
- `PROTOBUF_VERSION`: The version of Protobuf to use, for example [`3.0.0`]. Note: Changing this will not configure CMake to use a system version of Protobuf, it will configure CMake to download and try building that version.
- `CMAKE_TOOLCHAIN_FILE`: The path to a toolchain file for cross compilation.
- `BUILD_PARSERS`: Specify if the parsers should be built, for example [`ON`] | `OFF`. If turned OFF, CMake will try to find precompiled versions of the parser libraries to use in compiling samples. First in `${TRT_LIB_DIR}`, then on the system. If the build type is Debug, then it will prefer debug builds of the libraries before release versions if available.
- `BUILD_PLUGINS`: Specify if the plugins should be built, for example [`ON`] | `OFF`. If turned OFF, CMake will try to find a precompiled version of the plugin library to use in compiling samples. First in `${TRT_LIB_DIR}`, then on the system. If the build type is Debug, then it will prefer debug builds of the libraries before release versions if available.
- `BUILD_SAMPLES`: Specify if the samples should be built, for example [`ON`] | `OFF`.
- `GPU_ARCHS`: GPU (SM) architectures to target. By default we generate CUDA code for all major SMs. Specific SM versions can be specified here as a quoted space-separated list to reduce compilation time and binary size. Table of compute capabilities of NVIDIA GPUs can be found [here](https://developer.nvidia.com/cuda-gpus). Examples:
- NVidia A100: `-DGPU_ARCHS="80"`
- Tesla T4, GeForce RTX 2080: `-DGPU_ARCHS="75"`
- Titan V, Tesla V100: `-DGPU_ARCHS="70"`
- Multiple SMs: `-DGPU_ARCHS="80 75"`
- `TRT_PLATFORM_ID`: Bare-metal build (unlike containerized cross-compilation). Currently supported options: `x86_64` (default).
**Example: Native builds on Windows (x86) with cuda-12.8**
```bash
cd $TRT_OSSPATH
mkdir -p build
cd -p build
cmake .. -DTRT_LIB_DIR="$env:TRT_LIBPATH" -DCUDNN_ROOT_DIR="$env:CUDNN_PATH" -DTRT_OUT_DIR="$pwd\\out"
msbuild TensorRT.sln /property:Configuration=Release -m:$env:NUMBER_OF_PROCESSORS
```
> NOTE: The default CUDA version used by CMake is 12.8.0. To override this, for example to 11.8, append `-DCUDA_VERSION=11.8` to the cmake command.
- Required CMake build arguments are:
- `TRT_LIB_DIR`: Path to the TensorRT installation directory containing libraries.
- `TRT_OUT_DIR`: Output directory where generated build artifacts will be copied.
- Optional CMake build arguments:
- `CMAKE_BUILD_TYPE`: Specify if binaries generated are for release or debug (contain debug symbols). Values consists of [`Release`] | `Debug`
- `CUDA_VERSION`: The version of CUDA to target, for example [`11.7.1`].
- `CUDNN_VERSION`: The version of cuDNN to target, for example [`8.6`].
- `PROTOBUF_VERSION`: The version of Protobuf to use, for example [`3.0.0`]. Note: Changing this will not configure CMake to use a system version of Protobuf, it will configure CMake to download and try building that version.
- `CMAKE_TOOLCHAIN_FILE`: The path to a toolchain file for cross compilation.
- `BUILD_PARSERS`: Specify if the parsers should be built, for example [`ON`] | `OFF`. If turned OFF, CMake will try to find precompiled versions of the parser libraries to use in compiling samples. First in `${TRT_LIB_DIR}`, then on the system. If the build type is Debug, then it will prefer debug builds of the libraries before release versions if available.
- `BUILD_PLUGINS`: Specify if the plugins should be built, for example [`ON`] | `OFF`. If turned OFF, CMake will try to find a precompiled version of the plugin library to use in compiling samples. First in `${TRT_LIB_DIR}`, then on the system. If the build type is Debug, then it will prefer debug builds of the libraries before release versions if available.
- `BUILD_SAMPLES`: Specify if the samples should be built, for example [`ON`] | `OFF`.
- `GPU_ARCHS`: GPU (SM) architectures to target. By default we generate CUDA code for all major SMs. Specific SM versions can be specified here as a quoted space-separated list to reduce compilation time and binary size. Table of compute capabilities of NVIDIA GPUs can be found [here](https://developer.nvidia.com/cuda-gpus). Examples: - NVidia A100: `-DGPU_ARCHS="80"` - Tesla T4, GeForce RTX 2080: `-DGPU_ARCHS="75"` - Titan V, Tesla V100: `-DGPU_ARCHS="70"` - Multiple SMs: `-DGPU_ARCHS="80 75"`
- `TRT_PLATFORM_ID`: Bare-metal build (unlike containerized cross-compilation). Currently supported options: `x86_64` (default).
# References
## TensorRT Resources
* [TensorRT Developer Home](https://developer.nvidia.com/tensorrt)
* [TensorRT QuickStart Guide](https://docs.nvidia.com/deeplearning/tensorrt/quick-start-guide/index.html)
* [TensorRT Developer Guide](https://docs.nvidia.com/deeplearning/tensorrt/developer-guide/index.html)
* [TensorRT Sample Support Guide](https://docs.nvidia.com/deeplearning/tensorrt/sample-support-guide/index.html)
* [TensorRT ONNX Tools](https://docs.nvidia.com/deeplearning/tensorrt/index.html#tools)
* [TensorRT Discussion Forums](https://devtalk.nvidia.com/default/board/304/tensorrt/)
* [TensorRT Release Notes](https://docs.nvidia.com/deeplearning/tensorrt/release-notes/index.html)
- [TensorRT Developer Home](https://developer.nvidia.com/tensorrt)
- [TensorRT QuickStart Guide](https://docs.nvidia.com/deeplearning/tensorrt/quick-start-guide/index.html)
- [TensorRT Developer Guide](https://docs.nvidia.com/deeplearning/tensorrt/developer-guide/index.html)
- [TensorRT Sample Support Guide](https://docs.nvidia.com/deeplearning/tensorrt/sample-support-guide/index.html)
- [TensorRT ONNX Tools](https://docs.nvidia.com/deeplearning/tensorrt/index.html#tools)
- [TensorRT Discussion Forums](https://devtalk.nvidia.com/default/board/304/tensorrt/)
- [TensorRT Release Notes](https://docs.nvidia.com/deeplearning/tensorrt/release-notes/index.html)
## Known Issues
* Please refer to [TensorRT Release Notes](https://docs.nvidia.com/deeplearning/tensorrt/release-notes)
- Please refer to [TensorRT Release Notes](https://docs.nvidia.com/deeplearning/tensorrt/release-notes)
VERSION
+1 -1
View File
@@ -1 +1 @@
10.8.0.43
10.9.0.34
cmake/toolchains/cmake_aarch64_cross.toolchain
+2
View File
@@ -53,3 +53,5 @@ set(CMAKE_CUDA_COMPILER_FORCED TRUE)
set(CUDA_LIBS -L${CUDA_ROOT}/lib)
set(ADDITIONAL_PLATFORM_LIB_FLAGS ${CUDA_LIBS} -lcublas -lcudart -lstdc++ -lm)
link_directories(${CUDA_ROOT}/lib)
demo/BERT/README.md
+343 -311
View File
@@ -2,58 +2,56 @@
This subfolder of the BERT TensorFlow repository, tested and maintained by NVIDIA, provides scripts to perform high-performance inference using NVIDIA TensorRT.
## Table Of Contents
- [Model Overview](#model-overview)
* [Model Architecture](#model-architecture)
* [TensorRT Inference Pipeline](#tensorrt-inference-pipeline)
* [Version Info](#version-info)
- [Model Architecture](#model-architecture)
- [TensorRT Inference Pipeline](#tensorrt-inference-pipeline)
- [Version Info](#version-info)
- [Setup](#setup)
* [Requirements](#requirements)
- [Requirements](#requirements)
- [Quick Start Guide](#quick-start-guide)
* [(Optional) Trying a different configuration](#optional-trying-a-different-configuration)
- [(Optional) Trying a different configuration](#optional-trying-a-different-configuration)
- [Advanced](#advanced)
* [Scripts and sample code](#scripts-and-sample-code)
* [Command-line options](#command-line-options)
* [TensorRT inference process](#tensorrt-inference-process)
- [Scripts and sample code](#scripts-and-sample-code)
- [Command-line options](#command-line-options)
- [TensorRT inference process](#tensorrt-inference-process)
- [Accuracy](#accuracy)
* [Evaluating Post-Training-Quantization INT8 accuracy](#evaluating-ptq-post-training-quantization-int8-accuracy-using-the-squad-dataset)
* [Evaluating Quantization-Aware-Training INT8 accuracy](#evaluating-qat-quantization-aware-training-int8-accuracy-using-the-squad-dataset)
- [Evaluating Post-Training-Quantization INT8 accuracy](#evaluating-ptq-post-training-quantization-int8-accuracy-using-the-squad-dataset)
- [Evaluating Quantization-Aware-Training INT8 accuracy](#evaluating-qat-quantization-aware-training-int8-accuracy-using-the-squad-dataset)
- [Experimental](#experimental)
* [Variable sequence length](#variable-sequence-length)
* [Run command lines](#run-command-lines)
* [Sparsity with Quantization Aware Training](#sparsity-with-quantization-aware-training)
* [Megatron-LM for Question Answering](#megatron-lm-for-question-answering)
- [Variable sequence length](#variable-sequence-length)
- [Run command lines](#run-command-lines)
- [Sparsity with Quantization Aware Training](#sparsity-with-quantization-aware-training)
- [Megatron-LM for Question Answering](#megatron-lm-for-question-answering)
- [Performance](#performance)
* [Benchmarking](#benchmarking)
* [TensorRT inference benchmark](#tensorrt-inference-benchmark)
* [Results](#results)
* [Inference performance: NVIDIA A100](#inference-performance-nvidia-a100-40gb)
* [Inference performance: NVIDIA L4](#inference-performance-nvidia-l4)
* [Inference performance: NVIDIA L40S](#inference-performance-nvidia-l40s)
- [Benchmarking](#benchmarking)
- [TensorRT inference benchmark](#tensorrt-inference-benchmark)
- [Results](#results)
- [Inference performance: NVIDIA A100](#inference-performance-nvidia-a100-40gb)
- [Inference performance: NVIDIA L4](#inference-performance-nvidia-l4)
- [Inference performance: NVIDIA L40S](#inference-performance-nvidia-l40s)
## Model overview
BERT, or Bidirectional Encoder Representations from Transformers, is a new method of pre-training language representations which obtains state-of-the-art results on a wide array of Natural Language Processing (NLP) tasks. This model is based on the [BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding](https://arxiv.org/abs/1810.04805) paper. NVIDIA's BERT is an optimized version of [Google's official implementation](https://github.com/google-research/bert), leveraging mixed precision arithmetic and Tensor Cores for faster inference times while maintaining target accuracy.
Other publicly available implementations of BERT include:
1. [NVIDIA PyTorch](https://github.com/NVIDIA/DeepLearningExamples/tree/master/PyTorch/LanguageModeling/BERT)
2. [Hugging Face](https://github.com/huggingface/pytorch-pretrained-BERT)
3. [codertimo](https://github.com/codertimo/BERT-pytorch)
4. [gluon-nlp](https://github.com/dmlc/gluon-nlp/tree/master/scripts/bert)
5. [Google's official implementation](https://github.com/google-research/bert)
### Model architecture
BERT's model architecture is a multi-layer bidirectional Transformer encoder. Based on the model size, we have the following two default configurations of BERT:
| **Model** | **Hidden layers** | **Hidden unit size** | **Attention heads** | **Feed-forward filter size** | **Max sequence length** | **Parameters** |
|:---------:|:----------:|:----:|:---:|:--------:|:---:|:----:|
|BERT-Base |12 encoder| 768| 12|4 x 768|512|110M|
|BERT-Large|24 encoder|1024| 16|4 x 1024|512|330M|
| **Model** | **Hidden layers** | **Hidden unit size** | **Attention heads** | **Feed-forward filter size** | **Max sequence length** | **Parameters** |
| :--------: | :---------------: | :------------------: | :-----------------: | :--------------------------: | :---------------------: | :------------: |
| BERT-Base | 12 encoder | 768 | 12 | 4 x 768 | 512 | 110M |
| BERT-Large | 24 encoder | 1024 | 16 | 4 x 1024 | 512 | 330M |
Typically, the language model is followed by a few task-specific layers. The model used here includes layers for question answering.
@@ -72,11 +70,11 @@ Lastly, the tokens predicted by the model are projected back to the original tex
The following software version configuration has been tested:
|Software|Version|
|--------|-------|
|Python|>=3.8|
|TensorRT|10.7.0.23|
|CUDA|12.6|
| Software | Version |
| -------- | ------- |
| Python | >=3.8 |
| TensorRT | 10.9 |
| CUDA | 12.8 |
## Setup
@@ -86,102 +84,112 @@ The following section lists the requirements that you need to meet in order to r
This demo BERT application can be run within the TensorRT OSS build container. If running in a different environment, following packages are required.
* [NGC CLI](https://ngc.nvidia.com/setup/installers/cli) - for downloading BERT checkpoints from NGC.
* PyPI Packages:
* [pycuda](https://pypi.org/project/pycuda/) (tested v2019.1.2)
* [onnx](https://pypi.org/project/onnx) (tested v1.12.0)
* [tensorflow](https://pypi.org/project/tensorflow/) (tested v2.9.1)
* [torch](https://pypi.org/project/torch/) (tested v1.11.0)
* NVIDIA [Volta](https://www.nvidia.com/en-us/data-center/volta-gpu-architecture/), [Turing](https://www.nvidia.com/en-us/geforce/turing/) or [Ampere](https://www.nvidia.com/en-us/data-center/nvidia-ampere-gpu-architecture/) based GPU.
- [NGC CLI](https://ngc.nvidia.com/setup/installers/cli) - for downloading BERT checkpoints from NGC.
- PyPI Packages:
- [pycuda](https://pypi.org/project/pycuda/) (tested v2019.1.2)
- [onnx](https://pypi.org/project/onnx) (tested v1.12.0)
- [tensorflow](https://pypi.org/project/tensorflow/) (tested v2.9.1)
- [torch](https://pypi.org/project/torch/) (tested v1.11.0)
- NVIDIA [Volta](https://www.nvidia.com/en-us/data-center/volta-gpu-architecture/), [Turing](https://www.nvidia.com/en-us/geforce/turing/) or [Ampere](https://www.nvidia.com/en-us/data-center/nvidia-ampere-gpu-architecture/) based GPU.
## Quick Start Guide
1. Build and launch the container as described in [TensorRT OSS README](https://github.com/NVIDIA/TensorRT/blob/master/README.md).
**Note:** After this point, all commands should be run from within the container.
**Note:** After this point, all commands should be run from within the container.
2. Verify TensorRT installation by printing the version:
For example:
```bash
python3 -c "import tensorrt as trt; print(trt.__version__)"
```
```bash
python3 -c "import tensorrt as trt; print(trt.__version__)"
```
3. Download the SQuAD dataset and BERT checkpoints:
```bash
cd $TRT_OSSPATH/demo/BERT
```
Download SQuAD v1.1 training and dev dataset.
```bash
bash ./scripts/download_squad.sh
```
```bash
cd $TRT_OSSPATH/demo/BERT
```
Download Tensorflow checkpoints for BERT large model with sequence length 128, fine-tuned for SQuAD v2.0.
```bash
bash scripts/download_model.sh
```
Download SQuAD v1.1 training and dev dataset.
**Note:** Since the datasets and checkpoints are stored in the directory mounted from the host, they do *not* need to be downloaded each time the container is launched.
```bash
bash ./scripts/download_squad.sh
```
Download Tensorflow checkpoints for BERT large model with sequence length 128, fine-tuned for SQuAD v2.0.
```bash
bash scripts/download_model.sh
```
**Note:** Since the datasets and checkpoints are stored in the directory mounted from the host, they do _not_ need to be downloaded each time the container is launched.
**Warning:** In the event of encountering an error message stating, "Missing API key and missing Email Authentication. This command requires an API key or authentication via browser login", the recommended steps for resolution are as follows:
* Generate an API key by logging in https://ngc.nvidia.com/setup/api-key and copy the generated API key.
* Execute the command `ngc config set` in the docker and paste the copied API key into the prompt as directed.
- Generate an API key by logging in https://ngc.nvidia.com/setup/api-key and copy the generated API key.
- Execute the command `ngc config set` in the docker and paste the copied API key into the prompt as directed.
Completing these steps should resolve the error you encountered and allow the command to proceed successfully.
4. Build a TensorRT engine. To build an engine, run the `builder.py` script. For example:
```bash
mkdir -p engines && python3 builder.py -m models/fine-tuned/bert_tf_ckpt_large_qa_squad2_amp_128_v19.03.1/model.ckpt -o engines/bert_large_128.engine -b 1 -s 128 --fp16 -c models/fine-tuned/bert_tf_ckpt_large_qa_squad2_amp_128_v19.03.1
```
This will build an engine with a maximum batch size of 1 (`-b 1`), and sequence length of 128 (`-s 128`) using mixed precision (`--fp16`) using the BERT Large SQuAD v2 FP16 Sequence Length 128 checkpoint (`-c models/fine-tuned/bert_tf_ckpt_large_qa_squad2_amp_128_v19.03.1`).
```bash
mkdir -p engines && python3 builder.py -m models/fine-tuned/bert_tf_ckpt_large_qa_squad2_amp_128_v19.03.1/model.ckpt -o engines/bert_large_128.engine -b 1 -s 128 --fp16 -c models/fine-tuned/bert_tf_ckpt_large_qa_squad2_amp_128_v19.03.1
```
This will build an engine with a maximum batch size of 1 (`-b 1`), and sequence length of 128 (`-s 128`) using mixed precision (`--fp16`) using the BERT Large SQuAD v2 FP16 Sequence Length 128 checkpoint (`-c models/fine-tuned/bert_tf_ckpt_large_qa_squad2_amp_128_v19.03.1`).
5. Run inference. Two options are provided for running the model.
a. `inference.py` script
This script accepts a passage and question and then runs the engine to generate an answer.
For example:
```bash
python3 inference.py -e engines/bert_large_128.engine -p "TensorRT is a high performance deep learning inference platform that delivers low latency and high throughput for apps such as recommenders, speech and image/video on NVIDIA GPUs. It includes parsers to import models, and plugins to support novel ops and layers before applying optimizations for inference. Today NVIDIA is open-sourcing parsers and plugins in TensorRT so that the deep learning community can customize and extend these components to take advantage of powerful TensorRT optimizations for your apps." -q "What is TensorRT?" -v models/fine-tuned/bert_tf_ckpt_large_qa_squad2_amp_128_v19.03.1/vocab.txt
```
a. `inference.py` script
This script accepts a passage and question and then runs the engine to generate an answer.
For example:
b. `inference.ipynb` Jupyter Notebook
The Jupyter Notebook includes a passage and various example questions and allows you to interactively make modifications and see the outcome.
To launch the Jupyter Notebook from inside the container, run:
```bash
jupyter notebook --ip 0.0.0.0 inference.ipynb
```
Then, use your browser to open the link displayed. The link should look similar to: `http://127.0.0.1:8888/?token=<TOKEN>`
```bash
python3 inference.py -e engines/bert_large_128.engine -p "TensorRT is a high performance deep learning inference platform that delivers low latency and high throughput for apps such as recommenders, speech and image/video on NVIDIA GPUs. It includes parsers to import models, and plugins to support novel ops and layers before applying optimizations for inference. Today NVIDIA is open-sourcing parsers and plugins in TensorRT so that the deep learning community can customize and extend these components to take advantage of powerful TensorRT optimizations for your apps." -q "What is TensorRT?" -v models/fine-tuned/bert_tf_ckpt_large_qa_squad2_amp_128_v19.03.1/vocab.txt
```
b. `inference.ipynb` Jupyter Notebook
The Jupyter Notebook includes a passage and various example questions and allows you to interactively make modifications and see the outcome.
To launch the Jupyter Notebook from inside the container, run:
```bash
jupyter notebook --ip 0.0.0.0 inference.ipynb
```
Then, use your browser to open the link displayed. The link should look similar to: `http://127.0.0.1:8888/?token=<TOKEN>`
6. Run inference with CUDA Graph support.
A separate python `inference_c.py` script is provided to run inference with CUDA Graph support. This is necessary since CUDA Graph is only supported through CUDA C/C++ APIs, not pyCUDA. The `inference_c.py` script uses pybind11 to interface with C/C++ for CUDA graph capturing and launching. The cmdline interface is the same as `inference.py` except for an extra `--enable-graph` option.
A separate python `inference_c.py` script is provided to run inference with CUDA Graph support. This is necessary since CUDA Graph is only supported through CUDA C/C++ APIs, not pyCUDA. The `inference_c.py` script uses pybind11 to interface with C/C++ for CUDA graph capturing and launching. The cmdline interface is the same as `inference.py` except for an extra `--enable-graph` option.
```bash
mkdir -p build; pushd build
cmake .. -DPYTHON_EXECUTABLE=$(which python)
make -j
popd
python3 inference_c.py -e engines/bert_large_128.engine --enable-graph -p "TensorRT is a high performance deep learning inference platform that delivers low latency and high throughput for apps such as recommenders, speech and image/video on NVIDIA GPUs. It includes parsers to import models, and plugins to support novel ops and layers before applying optimizations for inference. Today NVIDIA is open-sourcing parsers and plugins in TensorRT so that the deep learning community can customize and extend these components to take advantage of powerful TensorRT optimizations for your apps." -q "What is TensorRT?" -v models/fine-tuned/bert_tf_ckpt_large_qa_squad2_amp_128_v19.03.1/vocab.txt
```
```bash
mkdir -p build; pushd build
cmake .. -DPYTHON_EXECUTABLE=$(which python)
make -j
popd
python3 inference_c.py -e engines/bert_large_128.engine --enable-graph -p "TensorRT is a high performance deep learning inference platform that delivers low latency and high throughput for apps such as recommenders, speech and image/video on NVIDIA GPUs. It includes parsers to import models, and plugins to support novel ops and layers before applying optimizations for inference. Today NVIDIA is open-sourcing parsers and plugins in TensorRT so that the deep learning community can customize and extend these components to take advantage of powerful TensorRT optimizations for your apps." -q "What is TensorRT?" -v models/fine-tuned/bert_tf_ckpt_large_qa_squad2_amp_128_v19.03.1/vocab.txt
```
A separate C/C++ inference benchmark executable `perf` (compiled from `perf.cpp`) is provided to run inference benchmarks with CUDA Graph. The cmdline interface is the same as `perf.py` except for an extra `--enable_graph` option.
```bash
build/perf -e engines/bert_large_128.engine -b 1 -s 128 -w 100 -i 1000 --enable_graph
```
A separate C/C++ inference benchmark executable `perf` (compiled from `perf.cpp`) is provided to run inference benchmarks with CUDA Graph. The cmdline interface is the same as `perf.py` except for an extra `--enable_graph` option.
```bash
build/perf -e engines/bert_large_128.engine -b 1 -s 128 -w 100 -i 1000 --enable_graph
```
### (Optional) Trying a different configuration
If you would like to run another configuration, you can manually download checkpoints using the included script. For example, run:
```bash
bash scripts/download_model.sh base
```
to download a BERT Base model instead of the default BERT Large model.
To view all available model options, run:
```bash
bash scripts/download_model.sh -h
```
@@ -212,6 +220,7 @@ Other folders included in the `root` directory are:
- `helpers` - Contains helpers for tokenization of inputs
The `infer_c/` folder contains all the necessary C/C++ files required for CUDA Graph support.
- `bert_infer.h` - Defines necessary data structures for running BERT inference
- `infer_c.cpp` - Defines C/C++ interface using pybind11 that can be plugged into `inference_c.py`
- `perf.cpp` - Runs inference benchmarks. It is equivalent to `perf.py`, with an extra option `--enable_graph` to enable CUDA Graph support.
@@ -228,32 +237,38 @@ If unspecified, `--use-deprecated-plugins` is used by default.
### TensorRT inference process
As mentioned in the [Quick Start Guide](#quick-start-guide), two options are provided for running inference:
1. The `inference.py` script which accepts a passage and a question and then runs the engine to generate an answer. Alternatively, this script can be used to run inference on the Squad dataset.
2. The `inference.ipynb` Jupyter Notebook which includes a passage and various example questions and allows you to interactively make modifications and see the outcome.
## Accuracy
### Evaluating PTQ (post-training quantization) Int8 Accuracy Using The SQuAD Dataset
1. Download Tensorflow checkpoints for a BERT Large FP16 SQuAD v2 model with a sequence length of 384:
```bash
bash scripts/download_model.sh large 384 v2
```
2. Build an engine:
2. Build an engine:
**Turing and Ampere GPUs**
```bash
# QKVToContextPlugin and SkipLayerNormPlugin supported with INT8 I/O. To enable, use -imh and -iln builder flags respectively.
mkdir -p engines && python3 builder.py -m models/fine-tuned/bert_tf_ckpt_large_qa_squad2_amp_384_v19.03.1/model.ckpt -o engines/bert_large_384_int8mix.engine -b 1 -s 384 --int8 --fp16 --strict -c models/fine-tuned/bert_tf_ckpt_large_qa_squad2_amp_384_v19.03.1 --squad-json ./squad/train-v1.1.json -v models/fine-tuned/bert_tf_ckpt_large_qa_squad2_amp_384_v19.03.1/vocab.txt --calib-num 100 -iln -imh
```
**Xavier GPU**
```bash
# Only supports SkipLayerNormPlugin running with INT8 I/O. Use -iln builder flag to enable.
mkdir -p engines && python3 builder.py -m models/fine-tuned/bert_tf_ckpt_large_qa_squad2_amp_384_v19.03.1/model.ckpt -o engines/bert_large_384_int8mix.engine -b 1 -s 384 --int8 --fp16 --strict -c models/fine-tuned/bert_tf_ckpt_large_qa_squad2_amp_384_v19.03.1 --squad-json ./squad/train-v1.1.json -v models/fine-tuned/bert_tf_ckpt_large_qa_squad2_amp_384_v19.03.1/vocab.txt --calib-num 100 -iln
```
**Volta GPU**
```bash
# No support for QKVToContextPlugin or SkipLayerNormPlugin running with INT8 I/O. Don't specify -imh or -iln in builder flags.
mkdir -p engines && python3 builder.py -m models/fine-tuned/bert_tf_ckpt_large_qa_squad2_amp_384_v19.03.1/model.ckpt -o engines/bert_large_384_int8mix.engine -b 1 -s 384 --int8 --fp16 --strict -c models/fine-tuned/bert_tf_ckpt_large_qa_squad2_amp_384_v19.03.1 --squad-json ./squad/train-v1.1.json -v models/fine-tuned/bert_tf_ckpt_large_qa_squad2_amp_384_v19.03.1/vocab.txt --calib-num 100
@@ -261,32 +276,38 @@ As mentioned in the [Quick Start Guide](#quick-start-guide), two options are pro
This will build an engine with a maximum batch size of 1 (`-b 1`), calibration dataset squad (`--squad-json ./squad/train-v1.1.json`), calibration sentences number 100 (`--calib-num 100`), and sequence length of 384 (`-s 384`) using INT8 mixed precision computation where possible (`--int8 --fp16 --strict`).
3. Run inference using the squad dataset, and evaluate the F1 score and exact match score:
3. Run inference using the squad dataset, and evaluate the F1 score and exact match score:
```bash
python3 inference.py -e engines/bert_large_384_int8mix.engine -s 384 -sq ./squad/dev-v1.1.json -v models/fine-tuned/bert_tf_ckpt_large_qa_squad2_amp_384_v19.03.1/vocab.txt -o ./predictions.json
python3 squad/evaluate-v1.1.py squad/dev-v1.1.json ./predictions.json 90
```
### Evaluating QAT (quantization aware training) Int8 Accuracy Using The SQuAD Dataset
1. Download checkpoint for BERT Large FP16 SQuAD v1.1 model with sequence length of 384:
```bash
bash scripts/download_model.sh pyt v1_1
```
2. Build an engine:
2. Build an engine:
**Turing and Ampere GPUs**
```bash
# QKVToContextPlugin and SkipLayerNormPlugin supported with INT8 I/O. To enable, use -imh and -iln builder flags respectively.
mkdir -p engines && python3 builder.py -o engines/bert_large_384_int8mix.engine -b 1 -s 384 --int8 --fp16 --strict -c models/fine-tuned/bert_tf_ckpt_large_qa_squad2_amp_384_v19.03.1 -v models/fine-tuned/bert_tf_ckpt_large_qa_squad2_amp_384_v19.03.1/vocab.txt -x models/fine-tuned/bert_pyt_onnx_large_qa_squad11_amp_fake_quant_v1/bert_large_v1_1_fake_quant.onnx -iln -imh
```
**Xavier GPU**
```bash
# Only supports SkipLayerNormPlugin running with INT8 I/O. Use -iln builder flag to enable.
mkdir -p engines && python3 builder.py -o engines/bert_large_384_int8mix.engine -b 1 -s 384 --int8 --fp16 --strict -c models/fine-tuned/bert_tf_ckpt_large_qa_squad2_amp_384_v19.03.1 -v models/fine-tuned/bert_tf_ckpt_large_qa_squad2_amp_384_v19.03.1/vocab.txt -x models/fine-tuned/bert_pyt_onnx_large_qa_squad11_amp_fake_quant_v1/bert_large_v1_1_fake_quant.onnx -iln
```
**Volta GPU**
```bash
# No support for QKVToContextPlugin or SkipLayerNormPlugin running with INT8 I/O. Don't specify -imh or -iln in builder flags.
mkdir -p engines && python3 builder.py -o engines/bert_large_384_int8mix.engine -b 1 -s 384 --int8 --fp16 --strict -c models/fine-tuned/bert_tf_ckpt_large_qa_squad2_amp_384_v19.03.1 -v models/fine-tuned/bert_tf_ckpt_large_qa_squad2_amp_384_v19.03.1/vocab.txt -x models/fine-tuned/bert_pyt_onnx_large_qa_squad11_amp_fake_quant_v1/bert_large_v1_1_fake_quant.onnx
@@ -294,42 +315,46 @@ As mentioned in the [Quick Start Guide](#quick-start-guide), two options are pro
This will build and engine with a maximum batch size of 1 (`-b 1`) and sequence length of 384 (`-s 384`) using INT8 mixed precision computation where possible (`--int8 --fp16 --strict`).
3. Run inference using the squad dataset, and evaluate the F1 score and exact match score:
3. Run inference using the squad dataset, and evaluate the F1 score and exact match score:
```bash
python3 inference.py -e engines/bert_large_384_int8mix.engine -s 384 -sq ./squad/dev-v1.1.json -v models/fine-tuned/bert_tf_ckpt_large_qa_squad2_amp_384_v19.03.1/vocab.txt -o ./predictions.json
python3 squad/evaluate-v1.1.py squad/dev-v1.1.json ./predictions.json 90
```
## Experimental
### Variable sequence length
In our prior implementation, we used inputs padded to max length along with corresponding input masks to handle variable sequence length inputs in a batch. The padding results in some wasted computations which can be avoided by handling variable sequence length inputs natively. Now we have a new approach called the variable sequence length method. By concatenating each input id into a single long input id, and concatenating each input segment id into a single long segment id, TensorRT can know the exact starts and ends by providing an extra sequence length buffer that contains the start and end positions of each sequence. Now we can eliminate the wasted computation in the input paddings.
Note this is an experimental feature because we only support Xavier+ GPUs, also there is neither FP32 support nor INT8 PTQ calibration.
1. Download checkpoint for BERT Large FP16 SQuAD v1.1 model with sequence length of 384:
```bash
bash scripts/download_model.sh pyt v1_1
```
2. Build an engine:
2. Build an engine:
**FP16 engine**
```bash
mkdir -p engines && python3 builder_varseqlen.py -x models/fine-tuned/bert_pyt_onnx_large_qa_squad11_amp_fake_quant_v1/bert_large_v1_1_fake_quant.onnx -o engines/bert_varseq_fp16.engine -b 1 -s 64 --fp16 -c models/fine-tuned/bert_tf_ckpt_large_qa_squad2_amp_384_v19.03.1 -v models/fine-tuned/bert_tf_ckpt_large_qa_squad2_amp_384_v19.03.1/vocab.txt
```
This will build and engine with a maximum batch size of 1 (`-b 1`) and sequence length of 64 (`-s 64`) using FP16 precision computation where possible (`--fp16`).
**INT8 engine**
```bash
mkdir -p engines && python3 builder_varseqlen.py -x models/fine-tuned/bert_pyt_onnx_large_qa_squad11_amp_fake_quant_v1/bert_large_v1_1_fake_quant.onnx -o engines/bert_varseq_int8.engine -b 1 -s 256 --int8 --fp16 -c models/fine-tuned/bert_tf_ckpt_large_qa_squad2_amp_384_v19.03.1 -v models/fine-tuned/bert_tf_ckpt_large_qa_squad2_amp_384_v19.03.1/vocab.txt
```
This will build and engine with a maximum batch size of 1 (`-b 1`) and sequence length of 256 (`-s 256`) using INT8 precision computation where possible (`--int8`).
3. Run inference
3. Run inference
Evaluate the F1 score and exact match score using the squad dataset:
@@ -344,7 +369,7 @@ Note this is an experimental feature because we only support Xavier+ GPUs, also
python3 inference_varseqlen.py -e engines/bert_varseq_int8.engine -p "TensorRT is a high performance deep learning inference platform that delivers low latency and high throughput for apps such as recommenders, speech and image/video on NVIDIA GPUs. It includes parsers to import models, and plugins to support novel ops and layers before applying optimizations for inference. Today NVIDIA is open-sourcing parsers and plugins in TensorRT so that the deep learning community can customize and extend these components to take advantage of powerful TensorRT optimizations for your apps." -q "What is TensorRT?" -v models/fine-tuned/bert_tf_ckpt_large_qa_squad2_amp_384_v19.03.1/vocab.txt -s 256
```
4. Collect performance data
4. Collect performance data
```bash
python3 perf_varseqlen.py -e engines/bert_varseq_int8.engine -b 1 -s 256
@@ -352,7 +377,7 @@ Note this is an experimental feature because we only support Xavier+ GPUs, also
This will collect performance data run use batch size 1 (`-b 1`) and sequence length of 256 (`-s 256`).
5. Collect performance data with CUDA graph enabled
5. Collect performance data with CUDA graph enabled
We can use the same `inference_c.py` and `build/perf` to collect performance data with cuda graph enabled. The command line is the same as run without variable sequence length.
@@ -373,10 +398,11 @@ The sparse weights are generated by finetuning with INT8 Quantization Aware Trai
**Build the TensorRT engine**:
Options specified:
* `--megatron` : assume Megatron style residuals instead of vanilla BERT.
* `--pickle` : specify a pickle file containing the PyTorch statedict corresponding to fine-tuned Megatron model.
* `-sp` : enable sparsity during engine optimization and treat the weights as sparse.
* `--int8 --il` : enable int8 tactics/plugins with interleaving.
- `--megatron` : assume Megatron style residuals instead of vanilla BERT.
- `--pickle` : specify a pickle file containing the PyTorch statedict corresponding to fine-tuned Megatron model.
- `-sp` : enable sparsity during engine optimization and treat the weights as sparse.
- `--int8 --il` : enable int8 tactics/plugins with interleaving.
```bash
bash ./scripts/download_model.sh 384 v1_1 # BERT-large model checkpoint fine-tuned for SQuAD 1.1
@@ -386,16 +412,20 @@ mkdir -p engines && python3 builder_varseqlen.py -c models/fine-tuned/bert_tf_ck
```
**Ask a question**:
```bash
python3 inference_varseqlen.py -e engines/megatron_large_seqlen384_int8qat_sparse.engine -p "TensorRT is a high performance deep learning inference platform that delivers low latency and high throughput for apps such as recommenders, speech and image/video on NVIDIA GPUs. It includes parsers to import models, and plugins to support novel ops and layers before applying optimizations for inference. Today NVIDIA is open-sourcing parsers and plugins in TensorRT so that the deep learning community can customize and extend these components to take advantage of powerful TensorRT optimizations for your apps." -q "What is TensorRT?" -v models/fine-tuned/bert_tf_ckpt_large_qa_squad2_amp_384_v19.03.1/vocab.txt -s 256
```
**Evaluate F1 score**:
```bash
python3 inference_varseqlen.py -e engines/megatron_large_seqlen384_int8qat_sparse.engine -s 384 -sq ./squad/dev-v1.1.json -v models/fine-tuned/bert_tf_ckpt_large_qa_squad2_amp_384_v19.03.1/vocab.txt -o ./predictions.json
python3 squad/evaluate-v1.1.py squad/dev-v1.1.json ./predictions.json 90
```
Expected output:
```
&&&& PASSED TensorRT BERT Squad Accuracy matches reference.
{"exact_match": 84.03973509933775, "f1": 90.88667129897755}
@@ -418,9 +448,11 @@ The inference benchmark is performed on a single GPU by the `inference_benchmark
**Note:** The time measurements do not include the time required to copy inputs to the device and copy outputs to the host.
To run the inference benchmark script, run:
```bash
bash scripts/inference_benchmark.sh --gpu <arch>
```
Options for `<arch>` are: 'Volta', 'Xavier', 'Turing', 'Ampere'
Note: Some of the configurations in the benchmark script require 16GB of GPU memory. On GPUs with smaller amounts of memory, parts of the benchmark may fail to run.
@@ -437,163 +469,163 @@ Results were obtained by running `scripts/inference_benchmark.sh --gpu Ampere` o
##### BERT base
| Sequence Length | Batch Size | INT8 Latency (ms) | | | FP16 Latency (ms) | | |
|-----------------|------------|-----------------|-----------------|---------|-----------------|-----------------|---------|
| | | 95th Percentile | 99th Percentile | Average | 95th Percentile | 99th Percentile | Average |
| 128 | 1 | 0.67 | 0.67 | 0.54 | 0.62 | 0.80 | 0.62 |
| 128 | 2 | 0.76 | 0.76 | 0.60 | 0.92 | 0.92 | 0.73 |
| 128 | 4 | 0.73 | 0.93 | 0.73 | 0.93 | 0.93 | 0.93 |
| 128 | 8 | 0.94 | 1.21 | 0.95 | 1.31 | 1.31 | 1.31 |
| 128 | 12 | 1.20 | 1.20 | 1.20 | 1.72 | 2.20 | 1.72 |
| 128 | 16 | 1.34 | 1.34 | 1.34 | 2.07 | 2.08 | 2.05 |
| 128 | 24 | 1.82 | 1.82 | 1.82 | 3.02 | 3.08 | 3.01 |
| 128 | 32 | 2.23 | 2.24 | 2.23 | 3.89 | 3.91 | 3.85 |
| 128 | 64 | 4.16 | 4.16 | 4.12 | 7.57 | 7.63 | 7.55 |
| 128 | 128 | 8.07 | 8.09 | 8.02 | 15.23 | 15.24 | 15.15 |
| 384 | 1 | 1.14 | 1.46 | 1.14 | 1.25 | 1.61 | 1.26 |
| 384 | 2 | 1.32 | 1.32 | 1.32 | 1.55 | 1.55 | 1.55 |
| 384 | 4 | 1.66 | 1.66 | 1.66 | 2.12 | 2.12 | 2.12 |
| 384 | 8 | 2.20 | 2.21 | 2.20 | 3.34 | 3.36 | 3.31 |
| 384 | 12 | 3.31 | 3.31 | 3.31 | 4.78 | 4.82 | 4.77 |
| 384 | 16 | 4.00 | 4.00 | 4.00 | 6.38 | 6.40 | 6.33 |
| 384 | 24 | 5.70 | 5.70 | 5.70 | 9.31 | 9.31 | 9.22 |
| 384 | 32 | 7.64 | 7.64 | 7.64 | 12.90 | 12.90 | 12.79 |
| 384 | 64 | 14.87 | 14.91 | 14.74 | 24.96 | 25.19 | 24.74 |
| 384 | 128 | 29.01 | 29.02 | 28.74 | 49.05 | 49.28 | 48.64 |
| Sequence Length | Batch Size | INT8 Latency (ms) | | | FP16 Latency (ms) | | |
| --------------- | ---------- | ----------------- | --------------- | ------- | ----------------- | --------------- | ------- |
| | | 95th Percentile | 99th Percentile | Average | 95th Percentile | 99th Percentile | Average |
| 128 | 1 | 0.67 | 0.67 | 0.54 | 0.62 | 0.80 | 0.62 |
| 128 | 2 | 0.76 | 0.76 | 0.60 | 0.92 | 0.92 | 0.73 |
| 128 | 4 | 0.73 | 0.93 | 0.73 | 0.93 | 0.93 | 0.93 |
| 128 | 8 | 0.94 | 1.21 | 0.95 | 1.31 | 1.31 | 1.31 |
| 128 | 12 | 1.20 | 1.20 | 1.20 | 1.72 | 2.20 | 1.72 |
| 128 | 16 | 1.34 | 1.34 | 1.34 | 2.07 | 2.08 | 2.05 |
| 128 | 24 | 1.82 | 1.82 | 1.82 | 3.02 | 3.08 | 3.01 |
| 128 | 32 | 2.23 | 2.24 | 2.23 | 3.89 | 3.91 | 3.85 |
| 128 | 64 | 4.16 | 4.16 | 4.12 | 7.57 | 7.63 | 7.55 |
| 128 | 128 | 8.07 | 8.09 | 8.02 | 15.23 | 15.24 | 15.15 |
| 384 | 1 | 1.14 | 1.46 | 1.14 | 1.25 | 1.61 | 1.26 |
| 384 | 2 | 1.32 | 1.32 | 1.32 | 1.55 | 1.55 | 1.55 |
| 384 | 4 | 1.66 | 1.66 | 1.66 | 2.12 | 2.12 | 2.12 |
| 384 | 8 | 2.20 | 2.21 | 2.20 | 3.34 | 3.36 | 3.31 |
| 384 | 12 | 3.31 | 3.31 | 3.31 | 4.78 | 4.82 | 4.77 |
| 384 | 16 | 4.00 | 4.00 | 4.00 | 6.38 | 6.40 | 6.33 |
| 384 | 24 | 5.70 | 5.70 | 5.70 | 9.31 | 9.31 | 9.22 |
| 384 | 32 | 7.64 | 7.64 | 7.64 | 12.90 | 12.90 | 12.79 |
| 384 | 64 | 14.87 | 14.91 | 14.74 | 24.96 | 25.19 | 24.74 |
| 384 | 128 | 29.01 | 29.02 | 28.74 | 49.05 | 49.28 | 48.64 |
##### BERT large
| Sequence Length | Batch Size | INT8 Latency (ms) | | | FP16 Latency (ms) | | |
|-----------------|------------|-----------------|-----------------|---------|-----------------|-----------------|---------|
| | | 95th Percentile | 99th Percentile | Average | 95th Percentile | 99th Percentile | Average |
| 128 | 1 | 1.23 | 1.23 | 1.23 | 1.54 | 1.55 | 1.54 |
| 128 | 2 | 1.42 | 1.42 | 1.42 | 1.82 | 2.02 | 1.82 |
| 128 | 4 | 1.79 | 1.79 | 1.78 | 2.52 | 2.53 | 2.52 |
| 128 | 8 | 2.64 | 2.65 | 2.64 | 3.93 | 3.94 | 3.89 |
| 128 | 12 | 3.11 | 3.11 | 3.11 | 5.03 | 5.07 | 5.00 |
| 128 | 16 | 4.09 | 4.09 | 4.08 | 6.93 | 6.94 | 6.86 |
| 128 | 24 | 5.28 | 5.28 | 5.27 | 9.70 | 9.70 | 9.65 |
| 128 | 32 | 7.00 | 7.02 | 6.95 | 12.95 | 12.96 | 12.83 |
| 128 | 64 | 12.85 | 12.89 | 12.74 | 24.85 | 25.06 | 24.63 |
| 128 | 128 | 25.07 | 25.08 | 24.99 | 49.15 | 49.42 | 48.69 |
| 384 | 1 | 2.55 | 2.55 | 2.55 | 2.96 | 2.96 | 2.96 |
| 384 | 2 | 3.03 | 3.03 | 3.03 | 3.90 | 3.90 | 3.89 |
| 384 | 4 | 4.01 | 4.01 | 4.01 | 5.73 | 5.79 | 5.67 |
| 384 | 8 | 7.16 | 7.16 | 7.16 | 11.12 | 11.16 | 11.01 |
| 384 | 12 | 9.14 | 9.14 | 9.13 | 15.31 | 15.45 | 15.27 |
| 384 | 16 | 12.28 | 12.28 | 12.28 | 20.99 | 20.99 | 20.92 |
| 384 | 24 | 17.67 | 17.72 | 17.57 | 30.75 | 31.03 | 30.66 |
| 384 | 32 | 23.29 | 23.31 | 23.06 | 41.01 | 41.26 | 40.61 |
| 384 | 64 | 44.96 | 45.30 | 44.83 | 79.97 | 80.27 | 79.26 |
| 384 | 128 | 87.99 | 88.02 | 87.69 | 156.51 | 156.99 | 155.47 |
| Sequence Length | Batch Size | INT8 Latency (ms) | | | FP16 Latency (ms) | | |
| --------------- | ---------- | ----------------- | --------------- | ------- | ----------------- | --------------- | ------- |
| | | 95th Percentile | 99th Percentile | Average | 95th Percentile | 99th Percentile | Average |
| 128 | 1 | 1.23 | 1.23 | 1.23 | 1.54 | 1.55 | 1.54 |
| 128 | 2 | 1.42 | 1.42 | 1.42 | 1.82 | 2.02 | 1.82 |
| 128 | 4 | 1.79 | 1.79 | 1.78 | 2.52 | 2.53 | 2.52 |
| 128 | 8 | 2.64 | 2.65 | 2.64 | 3.93 | 3.94 | 3.89 |
| 128 | 12 | 3.11 | 3.11 | 3.11 | 5.03 | 5.07 | 5.00 |
| 128 | 16 | 4.09 | 4.09 | 4.08 | 6.93 | 6.94 | 6.86 |
| 128 | 24 | 5.28 | 5.28 | 5.27 | 9.70 | 9.70 | 9.65 |
| 128 | 32 | 7.00 | 7.02 | 6.95 | 12.95 | 12.96 | 12.83 |
| 128 | 64 | 12.85 | 12.89 | 12.74 | 24.85 | 25.06 | 24.63 |
| 128 | 128 | 25.07 | 25.08 | 24.99 | 49.15 | 49.42 | 48.69 |
| 384 | 1 | 2.55 | 2.55 | 2.55 | 2.96 | 2.96 | 2.96 |
| 384 | 2 | 3.03 | 3.03 | 3.03 | 3.90 | 3.90 | 3.89 |
| 384 | 4 | 4.01 | 4.01 | 4.01 | 5.73 | 5.79 | 5.67 |
| 384 | 8 | 7.16 | 7.16 | 7.16 | 11.12 | 11.16 | 11.01 |
| 384 | 12 | 9.14 | 9.14 | 9.13 | 15.31 | 15.45 | 15.27 |
| 384 | 16 | 12.28 | 12.28 | 12.28 | 20.99 | 20.99 | 20.92 |
| 384 | 24 | 17.67 | 17.72 | 17.57 | 30.75 | 31.03 | 30.66 |
| 384 | 32 | 23.29 | 23.31 | 23.06 | 41.01 | 41.26 | 40.61 |
| 384 | 64 | 44.96 | 45.30 | 44.83 | 79.97 | 80.27 | 79.26 |
| 384 | 128 | 87.99 | 88.02 | 87.69 | 156.51 | 156.99 | 155.47 |
##### Megatron Large with Sparsity
| Sequence Length | Batch Size | INT8 QAT Latency (ms) | | |
|-----------------|------------|-----------------|-----------------|---------|
| | | 95th Percentile | 99th Percentile | Average |
| 128 | 1 | 1.13 | 1.44 | 1.14 |
| 128 | 2 | 1.37 | 1.37 | 1.37 |
| 128 | 4 | 1.78 | 1.78 | 1.77 |
| 128 | 8 | 2.45 | 2.46 | 2.45 |
| 128 | 12 | 3.11 | 3.12 | 3.10 |
| 128 | 16 | 3.91 | 3.91 | 3.90 |
| 128 | 24 | 4.89 | 4.89 | 4.88 |
| 128 | 32 | 6.96 | 6.97 | 6.91 |
| 128 | 64 | 11.64 | 11.65 | 11.63 |
| 128 | 128 | 21.82 | 21.83 | 21.69 |
| 384 | 1 | 1.69 | 1.69 | 1.69 |
| 384 | 2 | 2.21 | 2.22 | 2.21 |
| 384 | 4 | 3.63 | 3.63 | 3.62 |
| 384 | 8 | 5.72 | 5.72 | 5.71 |
| 384 | 12 | 8.38 | 8.39 | 8.37 |
| 384 | 16 | 10.35 | 10.35 | 10.34 |
| 384 | 24 | 14.49 | 14.49 | 14.48 |
| 384 | 32 | 18.75 | 18.81 | 18.73 |
| 384 | 64 | 36.38 | 36.41 | 36.11 |
| 384 | 128 | 69.82 | 69.95 | 69.34 |
| Sequence Length | Batch Size | INT8 QAT Latency (ms) | | |
| --------------- | ---------- | --------------------- | --------------- | ------- |
| | | 95th Percentile | 99th Percentile | Average |
| 128 | 1 | 1.13 | 1.44 | 1.14 |
| 128 | 2 | 1.37 | 1.37 | 1.37 |
| 128 | 4 | 1.78 | 1.78 | 1.77 |
| 128 | 8 | 2.45 | 2.46 | 2.45 |
| 128 | 12 | 3.11 | 3.12 | 3.10 |
| 128 | 16 | 3.91 | 3.91 | 3.90 |
| 128 | 24 | 4.89 | 4.89 | 4.88 |
| 128 | 32 | 6.96 | 6.97 | 6.91 |
| 128 | 64 | 11.64 | 11.65 | 11.63 |
| 128 | 128 | 21.82 | 21.83 | 21.69 |
| 384 | 1 | 1.69 | 1.69 | 1.69 |
| 384 | 2 | 2.21 | 2.22 | 2.21 |
| 384 | 4 | 3.63 | 3.63 | 3.62 |
| 384 | 8 | 5.72 | 5.72 | 5.71 |
| 384 | 12 | 8.38 | 8.39 | 8.37 |
| 384 | 16 | 10.35 | 10.35 | 10.34 |
| 384 | 24 | 14.49 | 14.49 | 14.48 |
| 384 | 32 | 18.75 | 18.81 | 18.73 |
| 384 | 64 | 36.38 | 36.41 | 36.11 |
| 384 | 128 | 69.82 | 69.95 | 69.34 |
### Inference Performance NVIDIA L4
#### Inference performance: NVIDIA A30
Results were obtained by running `scripts/inference_benchmark.sh --gpu Ampere` on NVIDIA L4.
Results were obtained by running `scripts/inference_benchmark.sh --gpu Ampere` on NVIDIA A30.
##### BERT base
| Sequence Length | Batch Size | INT8 Latency (ms) | | | FP16 Latency (ms) | | |
|-----------------|------------|-----------------|-----------------|---------|-----------------|-----------------|---------|
| | | 95th Percentile | 99th Percentile | Average | 95th Percentile | 99th Percentile | Average |
| 128 | 1 | 0.62 | 0.62 | 0.61 | 1.01 | 1.02 | 1.00 |
| 128 | 2 | 0.79 | 0.80 | 0.77 | 1.33 | 1.35 | 1.31 |
| 128 | 4 | 1.16 | 1.16 | 1.13 | 2.23 | 2.23 | 2.16 |
| 128 | 8 | 1.93 | 1.98 | 1.91 | 3.70 | 3.83 | 3.69 |
| 128 | 12 | 2.69 | 2.69 | 2.63 | 5.42 | 5.46 | 5.36 |
| 128 | 16 | 3.38 | 3.39 | 3.32 | 6.77 | 6.78 | 6.71 |
| 128 | 24 | 4.87 | 4.87 | 4.77 | 10.72 | 10.81 | 10.56 |
| 128 | 32 | 6.22 | 6.35 | 6.18 | 14.13 | 14.14 | 13.97 |
| 128 | 64 | 13.69 | 13.85 | 13.56 | 31.28 | 31.69 | 31.05 |
| 128 | 128 | 30.49 | 30.72 | 29.90 | 69.99 | 70.38 | 68.61 |
| 384 | 1 | 1.31 | 1.31 | 1.30 | 2.10 | 2.10 | 2.09 |
| 384 | 2 | 1.85 | 1.86 | 1.85 | 3.19 | 3.21 | 3.14 |
| 384 | 4 | 3.00 | 3.00 | 2.94 | 5.77 | 5.89 | 5.74 |
| 384 | 8 | 5.58 | 5.60 | 5.48 | 11.49 | 11.59 | 11.38 |
| 384 | 12 | 8.22 | 8.37 | 8.13 | 17.39 | 17.40 | 17.16 |
| 384 | 16 | 10.98 | 10.99 | 10.89 | 23.38 | 23.78 | 23.02 |
| 384 | 24 | 17.33 | 17.47 | 17.09 | 38.54 | 39.55 | 37.57 |
| 384 | 32 | 23.82 | 24.18 | 23.56 | 51.12 | 51.24 | 50.62 |
| 384 | 64 | 50.08 | 50.28 | 49.10 | 105.60 | 106.08 | 104.59 |
| 384 | 128 | 113.95 | 114.53 | 112.15 | 209.55 | 209.93 | 208.35 |
| Sequence Length | Batch Size | INT8 Latency (ms) | | | FP16 Latency (ms) | | |
| --------------- | ---------- | ----------------- | --------------- | ------- | ----------------- | --------------- | ------- |
| | | 95th Percentile | 99th Percentile | Average | 95th Percentile | 99th Percentile | Average |
| 128 | 1 | 0.62 | 0.62 | 0.61 | 1.01 | 1.02 | 1.00 |
| 128 | 2 | 0.79 | 0.80 | 0.77 | 1.33 | 1.35 | 1.31 |
| 128 | 4 | 1.16 | 1.16 | 1.13 | 2.23 | 2.23 | 2.16 |
| 128 | 8 | 1.93 | 1.98 | 1.91 | 3.70 | 3.83 | 3.69 |
| 128 | 12 | 2.69 | 2.69 | 2.63 | 5.42 | 5.46 | 5.36 |
| 128 | 16 | 3.38 | 3.39 | 3.32 | 6.77 | 6.78 | 6.71 |
| 128 | 24 | 4.87 | 4.87 | 4.77 | 10.72 | 10.81 | 10.56 |
| 128 | 32 | 6.22 | 6.35 | 6.18 | 14.13 | 14.14 | 13.97 |
| 128 | 64 | 13.69 | 13.85 | 13.56 | 31.28 | 31.69 | 31.05 |
| 128 | 128 | 30.49 | 30.72 | 29.90 | 69.99 | 70.38 | 68.61 |
| 384 | 1 | 1.31 | 1.31 | 1.30 | 2.10 | 2.10 | 2.09 |
| 384 | 2 | 1.85 | 1.86 | 1.85 | 3.19 | 3.21 | 3.14 |
| 384 | 4 | 3.00 | 3.00 | 2.94 | 5.77 | 5.89 | 5.74 |
| 384 | 8 | 5.58 | 5.60 | 5.48 | 11.49 | 11.59 | 11.38 |
| 384 | 12 | 8.22 | 8.37 | 8.13 | 17.39 | 17.40 | 17.16 |
| 384 | 16 | 10.98 | 10.99 | 10.89 | 23.38 | 23.78 | 23.02 |
| 384 | 24 | 17.33 | 17.47 | 17.09 | 38.54 | 39.55 | 37.57 |
| 384 | 32 | 23.82 | 24.18 | 23.56 | 51.12 | 51.24 | 50.62 |
| 384 | 64 | 50.08 | 50.28 | 49.10 | 105.60 | 106.08 | 104.59 |
| 384 | 128 | 113.95 | 114.53 | 112.15 | 209.55 | 209.93 | 208.35 |
##### BERT large
| Sequence Length | Batch Size | INT8 Latency (ms) | | | FP16 Latency (ms) | | |
|-----------------|------------|-----------------|-----------------|---------|-----------------|-----------------|---------|
| | | 95th Percentile | 99th Percentile | Average | 95th Percentile | 99th Percentile | Average |
| 128 | 1 | 1.80 | 1.80 | 1.78 | 3.12 | 3.12 | 3.10 |
| 128 | 2 | 2.51 | 2.52 | 2.45 | 4.36 | 4.38 | 4.34 |
| 128 | 4 | 3.70 | 3.72 | 3.60 | 6.91 | 6.94 | 6.82 |
| 128 | 8 | 6.40 | 6.41 | 6.30 | 12.81 | 12.96 | 12.76 |
| 128 | 12 | 8.53 | 8.60 | 8.36 | 18.79 | 18.96 | 18.43 |
| 128 | 16 | 11.25 | 11.34 | 11.18 | 25.61 | 25.85 | 25.34 |
| 128 | 24 | 16.25 | 16.28 | 16.01 | 36.21 | 36.22 | 35.95 |
| 128 | 32 | 21.65 | 21.68 | 21.37 | 49.57 | 49.76 | 49.17 |
| 128 | 64 | 44.98 | 45.44 | 44.57 | 107.87 | 108.20 | 106.77 |
| 128 | 128 | 93.97 | 94.63 | 93.01 | 216.03 | 216.54 | 214.69 |
| 384 | 1 | 3.47 | 3.48 | 3.45 | 6.64 | 6.75 | 6.43 |
| 384 | 2 | 5.57 | 5.58 | 5.46 | 10.63 | 10.65 | 10.49 |
| 384 | 4 | 9.79 | 9.93 | 9.62 | 20.78 | 21.19 | 20.38 |
| 384 | 8 | 18.38 | 18.39 | 18.22 | 39.85 | 40.17 | 38.38 |
| 384 | 12 | 26.50 | 26.74 | 26.39 | 61.30 | 61.76 | 59.94 |
| 384 | 16 | 37.19 | 37.48 | 36.70 | 81.72 | 82.15 | 80.66 |
| 384 | 24 | 55.13 | 55.69 | 54.64 | 131.37 | 131.61 | 130.29 |
| 384 | 32 | 76.86 | 77.41 | 75.98 | 166.22 | 166.56 | 165.16 |
| 384 | 64 | 165.08 | 165.56 | 163.82 | 344.18 | 344.61 | 342.97 |
| 384 | 128 | 334.73 | 335.97 | 332.16 | 670.67 | 671.67 | 668.80 |
| Sequence Length | Batch Size | INT8 Latency (ms) | | | FP16 Latency (ms) | | |
| --------------- | ---------- | ----------------- | --------------- | ------- | ----------------- | --------------- | ------- |
| | | 95th Percentile | 99th Percentile | Average | 95th Percentile | 99th Percentile | Average |
| 128 | 1 | 1.80 | 1.80 | 1.78 | 3.12 | 3.12 | 3.10 |
| 128 | 2 | 2.51 | 2.52 | 2.45 | 4.36 | 4.38 | 4.34 |
| 128 | 4 | 3.70 | 3.72 | 3.60 | 6.91 | 6.94 | 6.82 |
| 128 | 8 | 6.40 | 6.41 | 6.30 | 12.81 | 12.96 | 12.76 |
| 128 | 12 | 8.53 | 8.60 | 8.36 | 18.79 | 18.96 | 18.43 |
| 128 | 16 | 11.25 | 11.34 | 11.18 | 25.61 | 25.85 | 25.34 |
| 128 | 24 | 16.25 | 16.28 | 16.01 | 36.21 | 36.22 | 35.95 |
| 128 | 32 | 21.65 | 21.68 | 21.37 | 49.57 | 49.76 | 49.17 |
| 128 | 64 | 44.98 | 45.44 | 44.57 | 107.87 | 108.20 | 106.77 |
| 128 | 128 | 93.97 | 94.63 | 93.01 | 216.03 | 216.54 | 214.69 |
| 384 | 1 | 3.47 | 3.48 | 3.45 | 6.64 | 6.75 | 6.43 |
| 384 | 2 | 5.57 | 5.58 | 5.46 | 10.63 | 10.65 | 10.49 |
| 384 | 4 | 9.79 | 9.93 | 9.62 | 20.78 | 21.19 | 20.38 |
| 384 | 8 | 18.38 | 18.39 | 18.22 | 39.85 | 40.17 | 38.38 |
| 384 | 12 | 26.50 | 26.74 | 26.39 | 61.30 | 61.76 | 59.94 |
| 384 | 16 | 37.19 | 37.48 | 36.70 | 81.72 | 82.15 | 80.66 |
| 384 | 24 | 55.13 | 55.69 | 54.64 | 131.37 | 131.61 | 130.29 |
| 384 | 32 | 76.86 | 77.41 | 75.98 | 166.22 | 166.56 | 165.16 |
| 384 | 64 | 165.08 | 165.56 | 163.82 | 344.18 | 344.61 | 342.97 |
| 384 | 128 | 334.73 | 335.97 | 332.16 | 670.67 | 671.67 | 668.80 |
##### Megatron Large with Sparsity
| Sequence Length | Batch Size | INT8 QAT Latency (ms) | | |
|-----------------|------------|-----------------|-----------------|---------|
| | | 95th Percentile | 99th Percentile | Average |
| 128 | 1 | 1.51 | 1.51 | 1.49 |
| 128 | 2 | 2.07 | 2.09 | 2.03 |
| 128 | 4 | 2.98 | 3.02 | 2.92 |
| 128 | 8 | 5.06 | 5.07 | 5.05 |
| 128 | 12 | 6.70 | 6.77 | 6.63 |
| 128 | 16 | 8.81 | 8.82 | 8.74 |
| 128 | 24 | 13.18 | 13.19 | 13.09 |
| 128 | 32 | 17.43 | 17.44 | 17.34 |
| 128 | 64 | 36.26 | 36.70 | 35.86 |
| 128 | 128 | 79.70 | 79.88 | 79.06 |
| 384 | 1 | 2.80 | 2.81 | 2.75 |
| 384 | 2 | 4.21 | 4.21 | 4.15 |
| 384 | 4 | 7.64 | 7.66 | 7.53 |
| 384 | 8 | 14.96 | 14.98 | 14.83 |
| 384 | 12 | 21.62 | 21.66 | 21.46 |
| 384 | 16 | 28.40 | 28.57 | 28.31 |
| 384 | 24 | 45.11 | 45.45 | 44.78 |
| 384 | 32 | 60.86 | 61.08 | 59.88 |
| 384 | 64 | 126.53 | 126.80 | 126.06 |
| 384 | 128 | 255.35 | 256.27 | 253.63 |
| Sequence Length | Batch Size | INT8 QAT Latency (ms) | | |
| --------------- | ---------- | --------------------- | --------------- | ------- |
| | | 95th Percentile | 99th Percentile | Average |
| 128 | 1 | 1.51 | 1.51 | 1.49 |
| 128 | 2 | 2.07 | 2.09 | 2.03 |
| 128 | 4 | 2.98 | 3.02 | 2.92 |
| 128 | 8 | 5.06 | 5.07 | 5.05 |
| 128 | 12 | 6.70 | 6.77 | 6.63 |
| 128 | 16 | 8.81 | 8.82 | 8.74 |
| 128 | 24 | 13.18 | 13.19 | 13.09 |
| 128 | 32 | 17.43 | 17.44 | 17.34 |
| 128 | 64 | 36.26 | 36.70 | 35.86 |
| 128 | 128 | 79.70 | 79.88 | 79.06 |
| 384 | 1 | 2.80 | 2.81 | 2.75 |
| 384 | 2 | 4.21 | 4.21 | 4.15 |
| 384 | 4 | 7.64 | 7.66 | 7.53 |
| 384 | 8 | 14.96 | 14.98 | 14.83 |
| 384 | 12 | 21.62 | 21.66 | 21.46 |
| 384 | 16 | 28.40 | 28.57 | 28.31 |
| 384 | 24 | 45.11 | 45.45 | 44.78 |
| 384 | 32 | 60.86 | 61.08 | 59.88 |
| 384 | 64 | 126.53 | 126.80 | 126.06 |
| 384 | 128 | 255.35 | 256.27 | 253.63 |
### Inference Performance NVIDIA L40S
@@ -601,81 +633,81 @@ Results were obtained by running `scripts/inference_benchmark.sh --gpu Ampere` o
##### BERT base
| Sequence Length | Batch Size | INT8 Latency (ms) | | | FP16 Latency (ms) | | |
|-----------------|------------|-----------------|-----------------|---------|-----------------|-----------------|---------|
| | | 95th Percentile | 99th Percentile | Average | 95th Percentile | 99th Percentile | Average |
| 128 | 1 | 0.34 | 0.34 | 0.34 | 0.48 | 0.48 | 0.48 |
| 128 | 2 | 0.41 | 0.41 | 0.41 | 0.56 | 0.56 | 0.55 |
| 128 | 4 | 0.50 | 0.51 | 0.50 | 0.77 | 0.78 | 0.77 |
| 128 | 8 | 0.68 | 0.68 | 0.67 | 1.26 | 1.26 | 1.25 |
| 128 | 12 | 0.91 | 0.91 | 0.91 | 1.69 | 1.69 | 1.68 |
| 128 | 16 | 1.11 | 1.11 | 1.11 | 2.24 | 2.24 | 2.23 |
| 128 | 24 | 1.46 | 1.46 | 1.46 | 3.18 | 3.19 | 3.18 |
| 128 | 32 | 1.82 | 1.82 | 1.81 | 3.94 | 3.94 | 3.93 |
| 128 | 64 | 3.44 | 3.44 | 3.42 | 7.98 | 8.08 | 7.90 |
| 128 | 128 | 7.25 | 7.29 | 7.20 | 17.35 | 17.40 | 17.13 |
| 384 | 1 | 0.73 | 0.73 | 0.73 | 1.04 | 1.04 | 1.03 |
| 384 | 2 | 0.88 | 0.88 | 0.88 | 1.35 | 1.35 | 1.35 |
| 384 | 4 | 1.17 | 1.17 | 1.17 | 2.14 | 2.14 | 2.13 |
| 384 | 8 | 1.70 | 1.71 | 1.69 | 3.47 | 3.47 | 3.46 |
| 384 | 12 | 2.72 | 2.72 | 2.72 | 5.08 | 5.09 | 5.06 |
| 384 | 16 | 3.26 | 3.26 | 3.24 | 7.18 | 7.19 | 7.15 |
| 384 | 24 | 4.94 | 4.94 | 4.89 | 9.98 | 10.00 | 9.92 |
| 384 | 32 | 6.11 | 6.13 | 6.09 | 13.35 | 13.38 | 13.25 |
| 384 | 64 | 12.96 | 13.00 | 12.84 | 28.93 | 29.37 | 28.41 |
| 384 | 128 | 27.22 | 27.36 | 26.87 | 59.55 | 59.91 | 58.44 |
| Sequence Length | Batch Size | INT8 Latency (ms) | | | FP16 Latency (ms) | | |
| --------------- | ---------- | ----------------- | --------------- | ------- | ----------------- | --------------- | ------- |
| | | 95th Percentile | 99th Percentile | Average | 95th Percentile | 99th Percentile | Average |
| 128 | 1 | 0.34 | 0.34 | 0.34 | 0.48 | 0.48 | 0.48 |
| 128 | 2 | 0.41 | 0.41 | 0.41 | 0.56 | 0.56 | 0.55 |
| 128 | 4 | 0.50 | 0.51 | 0.50 | 0.77 | 0.78 | 0.77 |
| 128 | 8 | 0.68 | 0.68 | 0.67 | 1.26 | 1.26 | 1.25 |
| 128 | 12 | 0.91 | 0.91 | 0.91 | 1.69 | 1.69 | 1.68 |
| 128 | 16 | 1.11 | 1.11 | 1.11 | 2.24 | 2.24 | 2.23 |
| 128 | 24 | 1.46 | 1.46 | 1.46 | 3.18 | 3.19 | 3.18 |
| 128 | 32 | 1.82 | 1.82 | 1.81 | 3.94 | 3.94 | 3.93 |
| 128 | 64 | 3.44 | 3.44 | 3.42 | 7.98 | 8.08 | 7.90 |
| 128 | 128 | 7.25 | 7.29 | 7.20 | 17.35 | 17.40 | 17.13 |
| 384 | 1 | 0.73 | 0.73 | 0.73 | 1.04 | 1.04 | 1.03 |
| 384 | 2 | 0.88 | 0.88 | 0.88 | 1.35 | 1.35 | 1.35 |
| 384 | 4 | 1.17 | 1.17 | 1.17 | 2.14 | 2.14 | 2.13 |
| 384 | 8 | 1.70 | 1.71 | 1.69 | 3.47 | 3.47 | 3.46 |
| 384 | 12 | 2.72 | 2.72 | 2.72 | 5.08 | 5.09 | 5.06 |
| 384 | 16 | 3.26 | 3.26 | 3.24 | 7.18 | 7.19 | 7.15 |
| 384 | 24 | 4.94 | 4.94 | 4.89 | 9.98 | 10.00 | 9.92 |
| 384 | 32 | 6.11 | 6.13 | 6.09 | 13.35 | 13.38 | 13.25 |
| 384 | 64 | 12.96 | 13.00 | 12.84 | 28.93 | 29.37 | 28.41 |
| 384 | 128 | 27.22 | 27.36 | 26.87 | 59.55 | 59.91 | 58.44 |
##### BERT large
| Sequence Length | Batch Size | INT8 Latency (ms) | | | FP16 Latency (ms) | | |
|-----------------|------------|-----------------|-----------------|---------|-----------------|-----------------|---------|
| | | 95th Percentile | 99th Percentile | Average | 95th Percentile | 99th Percentile | Average |
| 128 | 1 | 0.89 | 0.89 | 0.89 | 1.30 | 1.30 | 1.30 |
| 128 | 2 | 0.98 | 0.98 | 0.98 | 1.45 | 1.46 | 1.45 |
| 128 | 4 | 1.35 | 1.35 | 1.34 | 2.32 | 2.32 | 2.31 |
| 128 | 8 | 1.93 | 1.95 | 1.92 | 3.59 | 3.60 | 3.58 |
| 128 | 12 | 2.73 | 2.73 | 2.72 | 5.70 | 5.71 | 5.63 |
| 128 | 16 | 3.19 | 3.21 | 3.17 | 6.48 | 6.49 | 6.45 |
| 128 | 24 | 4.50 | 4.53 | 4.48 | 9.89 | 9.90 | 9.81 |
| 128 | 32 | 5.66 | 5.68 | 5.62 | 12.26 | 12.30 | 12.16 |
| 128 | 64 | 11.42 | 11.43 | 11.30 | 27.40 | 27.60 | 27.16 |
| 128 | 128 | 24.68 | 24.70 | 24.36 | 61.49 | 61.76 | 60.81 |
| 384 | 1 | 1.68 | 1.68 | 1.68 | 2.73 | 2.73 | 2.73 |
| 384 | 2 | 2.28 | 2.28 | 2.27 | 3.83 | 3.83 | 3.82 |
| 384 | 4 | 3.28 | 3.28 | 3.26 | 6.26 | 6.26 | 6.24 |
| 384 | 8 | 4.97 | 4.98 | 4.95 | 10.32 | 10.33 | 10.30 |
| 384 | 12 | 7.89 | 7.89 | 7.86 | 17.49 | 17.50 | 17.43 |
| 384 | 16 | 9.47 | 9.49 | 9.44 | 21.50 | 21.62 | 21.24 |
| 384 | 24 | 14.64 | 14.66 | 14.54 | 33.26 | 33.30 | 33.01 |
| 384 | 32 | 19.20 | 19.37 | 18.97 | 44.56 | 44.69 | 43.95 |
| 384 | 64 | 42.15 | 42.38 | 41.56 | 98.89 | 99.28 | 97.19 |
| 384 | 128 | 84.15 | 84.40 | 83.34 | 196.98 | 197.83 | 194.18 |
| Sequence Length | Batch Size | INT8 Latency (ms) | | | FP16 Latency (ms) | | |
| --------------- | ---------- | ----------------- | --------------- | ------- | ----------------- | --------------- | ------- |
| | | 95th Percentile | 99th Percentile | Average | 95th Percentile | 99th Percentile | Average |
| 128 | 1 | 0.89 | 0.89 | 0.89 | 1.30 | 1.30 | 1.30 |
| 128 | 2 | 0.98 | 0.98 | 0.98 | 1.45 | 1.46 | 1.45 |
| 128 | 4 | 1.35 | 1.35 | 1.34 | 2.32 | 2.32 | 2.31 |
| 128 | 8 | 1.93 | 1.95 | 1.92 | 3.59 | 3.60 | 3.58 |
| 128 | 12 | 2.73 | 2.73 | 2.72 | 5.70 | 5.71 | 5.63 |
| 128 | 16 | 3.19 | 3.21 | 3.17 | 6.48 | 6.49 | 6.45 |
| 128 | 24 | 4.50 | 4.53 | 4.48 | 9.89 | 9.90 | 9.81 |
| 128 | 32 | 5.66 | 5.68 | 5.62 | 12.26 | 12.30 | 12.16 |
| 128 | 64 | 11.42 | 11.43 | 11.30 | 27.40 | 27.60 | 27.16 |
| 128 | 128 | 24.68 | 24.70 | 24.36 | 61.49 | 61.76 | 60.81 |
| 384 | 1 | 1.68 | 1.68 | 1.68 | 2.73 | 2.73 | 2.73 |
| 384 | 2 | 2.28 | 2.28 | 2.27 | 3.83 | 3.83 | 3.82 |
| 384 | 4 | 3.28 | 3.28 | 3.26 | 6.26 | 6.26 | 6.24 |
| 384 | 8 | 4.97 | 4.98 | 4.95 | 10.32 | 10.33 | 10.30 |
| 384 | 12 | 7.89 | 7.89 | 7.86 | 17.49 | 17.50 | 17.43 |
| 384 | 16 | 9.47 | 9.49 | 9.44 | 21.50 | 21.62 | 21.24 |
| 384 | 24 | 14.64 | 14.66 | 14.54 | 33.26 | 33.30 | 33.01 |
| 384 | 32 | 19.20 | 19.37 | 18.97 | 44.56 | 44.69 | 43.95 |
| 384 | 64 | 42.15 | 42.38 | 41.56 | 98.89 | 99.28 | 97.19 |
| 384 | 128 | 84.15 | 84.40 | 83.34 | 196.98 | 197.83 | 194.18 |
##### Megatron Large with Sparsity
| Sequence Length | Batch Size | INT8 QAT Latency (ms) | | |
|-----------------|------------|-----------------|-----------------|---------|
| | | 95th Percentile | 99th Percentile | Average |
| 128 | 1 | 0.76 | 0.76 | 0.76 |
| 128 | 2 | 0.90 | 0.90 | 0.90 |
| 128 | 4 | 1.13 | 1.13 | 1.13 |
| 128 | 8 | 1.71 | 1.71 | 1.71 |
| 128 | 12 | 2.26 | 2.26 | 2.25 |
| 128 | 16 | 2.72 | 2.73 | 2.72 |
| 128 | 24 | 4.44 | 4.45 | 4.43 |
| 128 | 32 | 5.07 | 5.11 | 5.04 |
| 128 | 64 | 10.06 | 10.09 | 9.97 |
| 128 | 128 | 20.42 | 20.46 | 20.30 |
| 384 | 1 | 1.13 | 1.13 | 1.13 |
| 384 | 2 | 1.63 | 1.65 | 1.62 |
| 384 | 4 | 2.52 | 2.53 | 2.51 |
| 384 | 8 | 4.93 | 4.94 | 4.90 |
| 384 | 12 | 6.47 | 6.47 | 6.45 |
| 384 | 16 | 8.41 | 8.42 | 8.36 |
| 384 | 24 | 12.52 | 12.53 | 12.44 |
| 384 | 32 | 16.66 | 16.72 | 16.57 |
| 384 | 64 | 34.12 | 34.22 | 33.81 |
| 384 | 128 | 71.98 | 72.13 | 71.52 |
| Sequence Length | Batch Size | INT8 QAT Latency (ms) | | |
| --------------- | ---------- | --------------------- | --------------- | ------- |
| | | 95th Percentile | 99th Percentile | Average |
| 128 | 1 | 0.76 | 0.76 | 0.76 |
| 128 | 2 | 0.90 | 0.90 | 0.90 |
| 128 | 4 | 1.13 | 1.13 | 1.13 |
| 128 | 8 | 1.71 | 1.71 | 1.71 |
| 128 | 12 | 2.26 | 2.26 | 2.25 |
| 128 | 16 | 2.72 | 2.73 | 2.72 |
| 128 | 24 | 4.44 | 4.45 | 4.43 |
| 128 | 32 | 5.07 | 5.11 | 5.04 |
| 128 | 64 | 10.06 | 10.09 | 9.97 |
| 128 | 128 | 20.42 | 20.46 | 20.30 |
| 384 | 1 | 1.13 | 1.13 | 1.13 |
| 384 | 2 | 1.63 | 1.65 | 1.62 |
| 384 | 4 | 2.52 | 2.53 | 2.51 |
| 384 | 8 | 4.93 | 4.94 | 4.90 |
| 384 | 12 | 6.47 | 6.47 | 6.45 |
| 384 | 16 | 8.41 | 8.42 | 8.36 |
| 384 | 24 | 12.52 | 12.53 | 12.44 |
| 384 | 32 | 16.66 | 16.72 | 16.57 |
| 384 | 64 | 34.12 | 34.22 | 33.81 |
| 384 | 128 | 71.98 | 72.13 | 71.52 |
## Hardware Platform Support
@@ -686,8 +718,8 @@ The compute capability of an NVIDIA GPU can be found out using the `nvidia-smi`
nvidia-smi --query-gpu=compute_cap --format=csv
```
Currently, this demo is supported on the following compute capabilities. This list is subject to change as new architectures are released.
- Volta architecture - 7.2, 7.5
- Ampere architecture - 8.0, 8.6, 8.7, 8.9
- Hopper architecture - 9.0 (since October 2022)
demo/DeBERTa/README.md
+3 -3
View File
@@ -75,7 +75,7 @@ Note that the performance gap between BERT's self-attention and DeBERTa's disent
## Environment Setup
It is recommended to use docker for reproducing the following steps. Follow the setup steps in TensorRT OSS [README](https://github.com/NVIDIA/TensorRT#setting-up-the-build-environment) to build and launch the container and build OSS:
**Example: Ubuntu 20.04 on x86-64 with cuda-12.5 (default)**
**Example: Ubuntu 20.04 on x86-64 with cuda-12.8 (default)**
```bash
# Download this TensorRT OSS repo
git clone -b main https://github.com/nvidia/TensorRT TensorRT
@@ -84,10 +84,10 @@ git submodule update --init --recursive
## at root of TensorRT OSS
# build container
./docker/build.sh --file docker/ubuntu-20.04.Dockerfile --tag tensorrt-ubuntu20.04-cuda12.5
./docker/build.sh --file docker/ubuntu-20.04.Dockerfile --tag tensorrt-ubuntu20.04-cuda12.8
# launch container
./docker/launch.sh --tag tensorrt-ubuntu20.04-cuda12.5 --gpus all
./docker/launch.sh --tag tensorrt-ubuntu20.04-cuda12.8 --gpus all
## now inside container
# build OSS (only required for pre-8.4.3 TensorRT versions)
demo/Diffusion/.gitignore
+5 -3
View File
@@ -1,4 +1,6 @@
__pycache__/
onnx/*.onnx
engine/*.plan
output/*.png
onnx/
engine/
output/
pytorch_model/
artifacts_cache/
demo/Diffusion/README.md
+40 -53
View File
@@ -7,7 +7,7 @@ This demo application ("demoDiffusion") showcases the acceleration of Stable Dif
### Clone the TensorRT OSS repository
```bash
git clone git@github.com:NVIDIA/TensorRT.git -b release/10.8 --single-branch
git clone git@github.com:NVIDIA/TensorRT.git -b release/10.9 --single-branch
cd TensorRT
```
@@ -42,13 +42,14 @@ pip3 install -r requirements.txt
```
> NOTE: demoDiffusion has been tested on systems with NVIDIA H100, A100, L40, T4, and RTX4090 GPUs, and the following software configuration.
```
diffusers 0.31.0
onnx 1.15.0
onnx-graphsurgeon 0.5.2
onnxruntime 1.16.3
polygraphy 0.49.9
tensorrt 10.8.0.43
tensorrt 10.9.0.34
tokenizers 0.13.3
torch 2.2.0
transformers 4.42.2
@@ -131,7 +132,6 @@ Examples:
Multiple ControlNet types can also be specified to combine the conditionings. While specifying multiple conditionings, controlnet scales should also be provided. The scales signify the importance of each conditioning in relation with the other. For example, to condition using `openpose` and `canny` with scales of 1.0 and 0.8 respectively, the arguments provided would be `--controlnet-type openpose canny` and `--controlnet-scale 1.0 0.8`. Note that the number of controlnet scales provided should match the number of controlnet types.
### Generate an image with Stable Diffusion XL guided by a single text prompt
Run the below command to generate an image with Stable Diffusion XL
@@ -146,6 +146,14 @@ The optional refiner model may be enabled by specifying `--enable-refiner` and s
python3 demo_txt2img_xl.py "a photo of an astronaut riding a horse on mars" --hf-token=$HF_TOKEN --version=xl-1.0 --enable-refiner --onnx-refiner-dir=onnx-refiner --engine-refiner-dir=engine-refiner
```
### Generate an image with Stable Diffusion XL with ControlNet guided by an image and a text prompt
```bash
python3 demo_controlnet.py "A beautiful bird with rainbow colors" --controlnet-type canny --hf-token=$HF_TOKEN --denoising-steps 20 --onnx-dir=onnx-cnet --engine-dir=engine-cnet --version xl-1.0
```
> NOTE: Currently only `--controlnet-type canny` is supported. `--input-image` must be a pre-processed image corresponding to `--controlnet-type canny`. If unspecified, a sample image will be downloaded.
### Generate an image guided by a text prompt, and using specified LoRA model weight updates
```bash
@@ -174,11 +182,15 @@ For step-by-step tutorials to run INT8 & FP8 inference on stable diffusion model
[LCM-LoRA](https://arxiv.org/abs/2311.05556) produces good quality images in 4 to 8 denoising steps instead of 30+ needed base model. Note that we use LCM scheduler and disable classifier-free-guidance by setting `--guidance-scale` to 0.
LoRA weights are fused into the ONNX and finalized TensorRT plan files in this example.
```bash
python3 demo_txt2img_xl.py "Einstein" --version xl-1.0 --lora-path "latent-consistency/lcm-lora-sdxl" --lora-weight 1.0 --onnx-dir onnx-sdxl-lcm-nocfg --engine-dir engine-sdxl-lcm-nocfg --denoising-steps 4 --scheduler LCM --guidance-scale 0.0
```
### Faster Text-to-Image using SDXL Turbo
Even faster image generation than LCM, producing coherent images in just 1 step. Note: SDXL Turbo works best for 512x512 resolution, EulerA scheduler and classifier-free-guidance disabled.
```bash
python3 demo_txt2img_xl.py "Einstein" --version xl-turbo --onnx-dir onnx-sdxl-turbo --engine-dir engine-sdxl-turbo --denoising-steps 1 --scheduler EulerA --guidance-scale 0.0 --width 512 --height 512
```
@@ -212,11 +224,11 @@ cd onnx-svd-xt-1-1 && git lfs pull && cd ..
```
SVD-XT-1.1 (25 frames at resolution 576x1024)
```bash
python3 demo_img2vid.py --version svd-xt-1.1 --onnx-dir onnx-svd-xt-1-1 --engine-dir engine-svd-xt-1-1 --hf-token=$HF_TOKEN
```
Run the command below to generate a video in FP8.
```bash
@@ -232,6 +244,7 @@ emb = emb.repeat_interleave(num_frames, dim=0)
```
You may also specify a custom conditioning image using `--input-image`:
```bash
python3 demo_img2vid.py --version svd-xt-1.1 --onnx-dir onnx-svd-xt-1-1 --engine-dir engine-svd-xt-1-1 --input-image https://www.hdcarwallpapers.com/walls/2018_chevrolet_camaro_zl1_nascar_race_car_2-HD.jpg --hf-token=$HF_TOKEN
```
@@ -241,11 +254,13 @@ NOTE: The min and max guidance scales are configured using --min-guidance-scale
### Generate an image guided by a text prompt using Stable Cascade
Run the below command to generate an image using Stable Cascade
```bash
python3 demo_stable_cascade.py --onnx-opset=16 "Anthropomorphic cat dressed as a pilot" --onnx-dir onnx-sc --engine-dir engine-sc
```
The lite versions of the models are also supported using the command below
```bash
python3 demo_stable_cascade.py --onnx-opset=16 "Anthropomorphic cat dressed as a pilot" --onnx-dir onnx-sc-lite --engine-dir engine-sc-lite --lite
```
@@ -258,61 +273,40 @@ python3 demo_stable_cascade.py --onnx-opset=16 "Anthropomorphic cat dressed as a
### Generating Images with Flux
#### Download Pre-exported Models (Recommended for GPUs with <48GB VRAM)
Install Git LFS:
```bash
sudo apt-get install git-lfs
```
Download ONNX models for the desired pipeline and precision:
```bash
# login to huggingface-cli using the $HF_TOKEN
git config --global credential.helper store # set the 'store' credential helper as default
huggingface-cli login --token $HF_TOKEN --add-to-git-credential
# Example for flux.1-dev BF16 pipeline. Models will be downloaded to ./onnx-flux-dev after the script is run.
./scripts/download_flux_onnx_models.sh --version "flux.1-dev" --precision "bf16"
# View supported configurations
./scripts/download_flux_onnx_models.sh --help
```
#### 1. Generate an Image from a Text Prompt
##### Run Flux.1-Dev
NOTE: Pass `--download-onnx-models` to avoid native ONNX export and download the ONNX models from (Black Forest Labs' collection)[https://huggingface.co/collections/black-forest-labs/flux1-onnx-679d06b7579583bd84c8ef83]. It is only supported for BF16, FP8, and FP4 pipelines.
```bash
# FP16 (requires >48GB VRAM for native export)
python3 demo_txt2img_flux.py "a beautiful photograph of Mt. Fuji during cherry blossom" --hf-token=$HF_TOKEN
# BF16
python3 demo_txt2img_flux.py "a beautiful photograph of Mt. Fuji during cherry blossom" --hf-token=$HF_TOKEN --bf16 --onnx-dir onnx-flux-dev/ --model-onnx-dirs=transformer:onnx-flux-dev/transformer_bf16/ --engine-dir engine-flux-dev/bf16
python3 demo_txt2img_flux.py "a beautiful photograph of Mt. Fuji during cherry blossom" --hf-token=$HF_TOKEN --bf16 --download-onnx-models
# FP8
python3 demo_txt2img_flux.py "a beautiful photograph of Mt. Fuji during cherry blossom" --hf-token=$HF_TOKEN --fp8 --onnx-dir onnx-flux-dev/ --model-onnx-dirs=transformer:onnx-flux-dev/transformer_fp8/ --engine-dir engine-flux-dev/fp8 --build-static-batch
python3 demo_txt2img_flux.py "a beautiful photograph of Mt. Fuji during cherry blossom" --hf-token=$HF_TOKEN --fp8 --download-onnx-models
# FP4
python3 demo_txt2img_flux.py "a beautiful photograph of Mt. Fuji during cherry blossom" --hf-token=$HF_TOKEN --fp4 --onnx-dir onnx-flux-dev/ --model-onnx-dirs=transformer:onnx-flux-dev/transformer_fp4/ --engine-dir engine-flux-dev/fp4
python3 demo_txt2img_flux.py "a beautiful photograph of Mt. Fuji during cherry blossom" --hf-token=$HF_TOKEN --fp4 --download-onnx-models
```
##### Run Flux.1-Schnell
```bash
# FP16 (requires >48GB VRAM for native export)
python3 demo_txt2img_flux.py "a beautiful photograph of Mt. Fuji during cherry blossom" --hf-token=$HF_TOKEN --version="flux.1-schnell"
# BF16
python3 demo_txt2img_flux.py "a beautiful photograph of Mt. Fuji during cherry blossom" --hf-token=$HF_TOKEN --version="flux.1-schnell" --bf16 --onnx-dir onnx-flux-schnell/ --model-onnx-dirs=transformer:onnx-flux-schnell/transformer_bf16 --engine-dir engine-flux-schnell/bf16
python3 demo_txt2img_flux.py "a beautiful photograph of Mt. Fuji during cherry blossom" --hf-token=$HF_TOKEN --version="flux.1-schnell" --bf16 --download-onnx-models
# FP8
python3 demo_txt2img_flux.py "a beautiful photograph of Mt. Fuji during cherry blossom" --hf-token=$HF_TOKEN --version="flux.1-schnell" --fp8 --onnx-dir onnx-flux-schnell/ --model-onnx-dirs=transformer:onnx-flux-schnell/transformer_fp8 --engine-dir engine-flux-schnell/fp8 --build-static-batch
python3 demo_txt2img_flux.py "a beautiful photograph of Mt. Fuji during cherry blossom" --hf-token=$HF_TOKEN --version="flux.1-schnell" --fp8 --download-onnx-models
# FP4
python3 demo_txt2img_flux.py "a beautiful photograph of Mt. Fuji during cherry blossom" --hf-token=$HF_TOKEN --version="flux.1-schnell" --fp4 --onnx-dir onnx-flux-schnell/ --model-onnx-dirs=transformer:onnx-flux-schnell/transformer_fp4 --engine-dir engine-flux-schnell/fp4
python3 demo_txt2img_flux.py "a beautiful photograph of Mt. Fuji during cherry blossom" --hf-token=$HF_TOKEN --version="flux.1-schnell" --fp4 --download-onnx-models
```
---
@@ -351,32 +345,32 @@ You can use the `--calibraton-dataset` flag to specify the path, which is set to
```bash
# BF16
python3 demo_img2img_flux.py "A robot made of exotic candies and chocolates of different kinds. The background is filled with confetti and celebratory gifts." --version="flux.1-dev-depth" --hf-token=$HF_TOKEN --guidance-scale 10 --control-image robot.png --bf16 --denoising-steps 30 --onnx-dir onnx-flux-dev-depth/ --model-onnx-dirs=transformer:onnx-flux-dev-depth/transformer_bf16 --engine-dir engine-flux-dev-depth/bf16
python3 demo_img2img_flux.py "A robot made of exotic candies and chocolates of different kinds. The background is filled with confetti and celebratory gifts." --version="flux.1-dev-depth" --hf-token=$HF_TOKEN --guidance-scale 10 --control-image robot.png --bf16 --denoising-steps 30 --download-onnx-models
# FP8 using pre-exported ONNX models
python3 demo_img2img_flux.py "A robot made of exotic candies" --version="flux.1-dev-depth" --hf-token=$HF_TOKEN --guidance-scale 10 --control-image robot.png --fp8 --denoising-steps 30 --onnx-dir onnx-flux-dev-depth/ --model-onnx-dirs=transformer:onnx-flux-dev-depth/transformer_fp8 --engine-dir engine-flux-dev-depth/fp8 --build-static-batch
python3 demo_img2img_flux.py "A robot made of exotic candies" --version="flux.1-dev-depth" --hf-token=$HF_TOKEN --guidance-scale 10 --control-image robot.png --fp8 --denoising-steps 30 --download-onnx-models --build-static-batch
# FP8 using native ONNX export
rm -rf onnx/* engine/* && python3 demo_img2img_flux.py "A robot made of exotic candies" --version="flux.1-dev-depth" --hf-token=$HF_TOKEN --guidance-scale 10 --control-image robot.png --fp8 --denoising-steps 30
# FP4
python3 demo_img2img_flux.py "A robot made of exotic candies" --version="flux.1-dev-depth" --hf-token=$HF_TOKEN --guidance-scale 10 --control-image robot.png --fp4 --denoising-steps 30 --onnx-dir onnx-flux-dev-depth/ --model-onnx-dirs=transformer:onnx-flux-dev-depth/transformer_fp4 --engine-dir engine-flux-dev-depth/fp4
python3 demo_img2img_flux.py "A robot made of exotic candies" --version="flux.1-dev-depth" --hf-token=$HF_TOKEN --guidance-scale 10 --control-image robot.png --fp4 --denoising-steps 30 --download-onnx-models
```
##### Canny ControlNet
```bash
# BF16
python3 demo_img2img_flux.py "a robot made out of gold" --version="flux.1-dev-canny" --hf-token=$HF_TOKEN --guidance-scale 30 --control-image robot.png --bf16 --onnx-dir onnx-flux-dev-canny/ --model-onnx-dirs=transformer:onnx-flux-dev-canny/transformer_bf16 --engine-dir engine-flux-dev-canny/bf16
python3 demo_img2img_flux.py "a robot made out of gold" --version="flux.1-dev-canny" --hf-token=$HF_TOKEN --guidance-scale 30 --control-image robot.png --bf16 --denoising-steps 30 --download-onnx-models
# FP8 using pre-exported ONNX models
python3 demo_img2img_flux.py "a robot made out of gold" --version="flux.1-dev-canny" --hf-token=$HF_TOKEN --guidance-scale 30 --control-image robot.png --fp8 --onnx-dir onnx-flux-dev-canny/ --model-onnx-dirs=transformer:onnx-flux-dev-canny/transformer_fp8 --engine-dir engine-flux-dev-canny/fp8 --build-static-batch
python3 demo_img2img_flux.py "a robot made out of gold" --version="flux.1-dev-canny" --hf-token=$HF_TOKEN --guidance-scale 30 --control-image robot.png --fp8 --denoising-steps 30 --download-onnx-models --build-static-batch
# FP8 using native ONNX export
rm -rf onnx/* engine/* && python3 demo_img2img_flux.py "a robot made out of gold" --version="flux.1-dev-canny" --hf-token=$HF_TOKEN --guidance-scale 30 --control-image robot.png --fp8 --calibration-dataset {custom/dataset/path}
rm -rf onnx/* engine/* && python3 demo_img2img_flux.py "a robot made out of gold" --version="flux.1-dev-canny" --hf-token=$HF_TOKEN --guidance-scale 30 --control-image robot.png --fp8 --denoising-steps 30 --calibration-dataset {custom/dataset/path}
# FP4
python3 demo_img2img_flux.py "a robot made out of gold" --version="flux.1-dev-canny" --hf-token=$HF_TOKEN --guidance-scale 30 --control-image robot.png --fp4 --onnx-dir onnx-flux-dev-canny/ --model-onnx-dirs=transformer:onnx-flux-dev-canny/transformer_fp4 --engine-dir engine-flux-dev-canny/fp4
python3 demo_img2img_flux.py "a robot made out of gold" --version="flux.1-dev-canny" --hf-token=$HF_TOKEN --guidance-scale 30 --control-image robot.png --fp4 --denoising-steps 30 --download-onnx-models
```
---
@@ -398,7 +392,7 @@ python3 demo_txt2img_flux.py "a beautiful photograph of Mt. Fuji during cherry b
- `--low-vram`: Enables model-offloading for reduced VRAM usage.
- `--ws`: Enables weight streaming in TensorRT engines.
- `--t5-ws-percentage` and `--transformer-ws-percentage`: Set runtime weight streaming budgets.
- `--build-static-batch`: Build all engines using static dimensions to lower the required activation memory. This will limit inference to the specified spatial dimensions.
- `--build-static-batch`: Build all engines using static batch sizes to lower the required activation memory. This will limit supported batch size of these engines for inference to the value specified by `--batch-size`.
##### FLUX VRAM Requirements Table
@@ -413,22 +407,15 @@ Memory usage captured below excludes the ONNX export step, and assumes use of th
NOTE: The FP8 and FP4 Pipelines are supported on Hopper/Ada/Blackwell devices only. The FP4 pipeline is most performant on Blackwell devices.
### Use separate directories for individual ONNX models
The directories specified in `--model-onnx-dirs` will override the directory set in `--onnx-dir`. Unspecified models will continue to use the directory set in `--onnx-dir`.
Suppose the model storage locations are as following:
* transformer model ONNX files are saved at `./onnx_folder_1/transformer` and `./onnx_folder_1/transformer.opt`.
* vae model ONNX files are saved in `./onnx_folder_2/vae` and `./onnx_folder_2/vae.opt`.
* Other models (t5 and clip) are still under `./onnx/`.
### Specify Custom Paths for ONNX models and TensorRT engines (FLUX only)
The corresponding command to run the pipeline:
```bash
python3 demo_txt2img_flux.py "a beautiful photograph of Mt. Fuji during cherry blossom" --hf-token=$HF_TOKEN --onnx-dir=onnx --model-onnx-dirs=transformer:onnx_folder_1,vae:onnx_folder_2
```
Custom override paths to pre-exported ONNX model files can be provided using `--custom-onnx-paths`. These ONNX models are directly used to build TRT engines without further optimization on the ONNX graphs. Paths should be a comma-separated list of <model_name>:<path> pairs. For example: `--custom-onnx-paths=transformer:/path/to/transformer.onnx,vae:/path/to/vae.onnx`. Call <PipelineClass>.get_model_names(...) for the list of supported model names.
Custom override paths to pre-built engine files can be provided using `--custom-engine-paths`. Paths should be a comma-separated list of <model_name>:<path> pairs. For example: `--custom-onnx-paths=transformer:/path/to/transformer.plan,vae:/path/to/vae.plan`.
## Configuration options
- Noise scheduler can be set using `--scheduler <scheduler>`. Note: not all schedulers are available for every version.
- To accelerate engine building time use `--timing-cache <path to cache file>`. The cache file will be created if it does not already exist. Note that performance may degrade if cache files are used across multiple GPU targets. It is recommended to use timing caches only during development. To achieve the best perfromance in deployment, please build engines without timing cache.
- Specify new directories for storing onnx and engine files when switching between versions, LoRAs, ControlNets, etc. This can be done using `--onnx-dir <new onnx dir>` and `--engine-dir <new engine dir>`.
- Inference performance can be improved by enabling [CUDA graphs](https://docs.nvidia.com/cuda/cuda-c-programming-guide/index.html#cuda-graphs) using `--use-cuda-graph`. Enabling CUDA graphs requires fixed input shapes, so this flag must be combined with `--build-static-batch` and cannot be combined with `--build-dynamic-shape`.
demo/Diffusion/calibration-images/rocket.png → demo/Diffusion/calibration_data/calibration-images/rocket.png
View File
Side by Side Swipe Overlay

Before

Width:  |  Height:  |  Size: 482 KiB

After

Width:  |  Height:  |  Size: 482 KiB

demo/Diffusion/calibration-prompts.txt → demo/Diffusion/calibration_data/calibration-prompts.txt
View File
demo/Diffusion/demo_controlnet.py
+53 -19
View File
@@ -22,12 +22,14 @@ import torch
from cuda import cudart
from PIL import Image
from stable_diffusion_pipeline import StableDiffusionPipeline
from utilities import PIPELINE_TYPE, TRT_LOGGER, add_arguments, download_image, process_pipeline_args
from demo_diffusion import dd_argparse
from demo_diffusion import image as image_module
from demo_diffusion import pipeline as pipeline_module
def parseArgs():
parser = argparse.ArgumentParser(description="Options for Stable Diffusion ControlNet Demo", conflict_handler='resolve')
parser = add_arguments(parser)
parser = dd_argparse.add_arguments(parser)
parser.add_argument('--scheduler', type=str, default="UniPC", choices=["DDIM", "DPM", "EulerA", "LMSD", "PNDM", "UniPC"], help="Scheduler for diffusion process")
parser.add_argument('--input-image', nargs = '+', type=str, default=[], help="Path to the input image/images already prepared for ControlNet modality. For example: canny edged image for canny ControlNet, not just regular rgb image")
parser.add_argument('--controlnet-type', nargs='+', type=str, default=["canny"], help="Controlnet type, can be `None`, `str` or `str` list from ['canny', 'depth', 'hed', 'mlsd', 'normal', 'openpose', 'scribble', 'seg']")
@@ -41,15 +43,15 @@ if __name__ == "__main__":
# Controlnet configuration
if not isinstance(args.controlnet_type, list):
raise ValueError(f"`--controlnet-type` must be of type `str` or `str` list, but is {type(args.controlnet_type)}")
# Controlnet configuration
if not isinstance(args.controlnet_scale, list):
raise ValueError(f"`--controlnet-scale`` must be of type `float` or `float` list, but is {type(args.controlnet_scale)}")
# Check number of ControlNets to ControlNet scales
if len(args.controlnet_type) != len(args.controlnet_scale):
raise ValueError(f"Numbers of ControlNets {len(args.controlnet_type)} should be equal to number of ControlNet scales {len(args.controlnet_scale)}.")
# Convert controlnet scales to tensor
controlnet_scale = torch.FloatTensor(args.controlnet_scale)
@@ -61,48 +63,80 @@ if __name__ == "__main__":
else:
for controlnet in args.controlnet_type:
if controlnet == "canny":
canny_image = download_image("https://hf.co/datasets/huggingface/documentation-images/resolve/main/diffusers/input_image_vermeer.png")
canny_image = controlnet_aux.CannyDetector()(canny_image)
if args.version == "xl-1.0":
canny_image = image_module.download_image(
"https://huggingface.co/diffusers/controlnet-canny-sdxl-1.0/resolve/main/out_bird.png"
)
# "out_bird.png" has 5 images combined in a row. We pick the first image which is the input image.
canny_image = canny_image.crop((0, 0, canny_image.width / 5, canny_image.height))
elif args.version == "1.5":
canny_image = image_module.download_image(
"https://hf.co/datasets/huggingface/documentation-images/resolve/main/diffusers/input_image_vermeer.png"
)
canny_image = controlnet_aux.CannyDetector()(canny_image)
else:
raise ValueError(
f"This demo supports ControlNets for v1.5 and SDXL base pipelines only. Version provided: {args.version}"
)
input_images.append(canny_image.resize((args.height, args.width)))
elif controlnet == "normal":
normal_image = download_image("https://huggingface.co/lllyasviel/sd-controlnet-normal/resolve/main/images/toy.png")
normal_image = image_module.download_image(
"https://huggingface.co/lllyasviel/sd-controlnet-normal/resolve/main/images/toy.png"
)
normal_image = controlnet_aux.NormalBaeDetector.from_pretrained("lllyasviel/Annotators")(normal_image)
input_images.append(normal_image.resize((args.height, args.width)))
elif controlnet == "depth":
depth_image = download_image("https://huggingface.co/lllyasviel/sd-controlnet-depth/resolve/main/images/stormtrooper.png")
depth_image = image_module.download_image(
"https://huggingface.co/lllyasviel/sd-controlnet-depth/resolve/main/images/stormtrooper.png"
)
depth_image = controlnet_aux.LeresDetector.from_pretrained("lllyasviel/Annotators")(depth_image)
input_images.append(depth_image.resize((args.height, args.width)))
elif controlnet == "hed":
hed_image = download_image("https://huggingface.co/lllyasviel/sd-controlnet-hed/resolve/main/images/man.png")
hed_image = image_module.download_image(
"https://huggingface.co/lllyasviel/sd-controlnet-hed/resolve/main/images/man.png"
)
hed_image = controlnet_aux.HEDdetector.from_pretrained("lllyasviel/Annotators")(hed_image)
input_images.append(hed_image.resize((args.height, args.width)))
elif controlnet == "mlsd":
mlsd_image = download_image("https://huggingface.co/lllyasviel/sd-controlnet-mlsd/resolve/main/images/room.png")
mlsd_image = image_module.download_image(
"https://huggingface.co/lllyasviel/sd-controlnet-mlsd/resolve/main/images/room.png"
)
mlsd_image = controlnet_aux.MLSDdetector.from_pretrained("lllyasviel/Annotators")(mlsd_image)
input_images.append(mlsd_image.resize((args.height, args.width)))
elif controlnet == "openpose":
openpose_image = download_image("https://huggingface.co/lllyasviel/sd-controlnet-openpose/resolve/main/images/pose.png")
openpose_image = image_module.download_image(
"https://huggingface.co/lllyasviel/sd-controlnet-openpose/resolve/main/images/pose.png"
)
openpose_image = controlnet_aux.OpenposeDetector.from_pretrained("lllyasviel/Annotators")(openpose_image)
input_images.append(openpose_image.resize((args.height, args.width)))
elif controlnet == "scribble":
scribble_image = download_image("https://huggingface.co/lllyasviel/sd-controlnet-scribble/resolve/main/images/bag.png")
scribble_image = image_module.download_image(
"https://huggingface.co/lllyasviel/sd-controlnet-scribble/resolve/main/images/bag.png"
)
scribble_image = controlnet_aux.HEDdetector.from_pretrained("lllyasviel/Annotators")(scribble_image, scribble=True)
input_images.append(scribble_image.resize((args.height, args.width)))
elif controlnet == "seg":
seg_image = download_image("https://huggingface.co/lllyasviel/sd-controlnet-seg/resolve/main/images/house.png")
seg_image = image_module.download_image(
"https://huggingface.co/lllyasviel/sd-controlnet-seg/resolve/main/images/house.png"
)
seg_image = controlnet_aux.SamDetector.from_pretrained("ybelkada/segment-anything", subfolder="checkpoints")(seg_image)
input_images.append(seg_image.resize((args.height, args.width)))
else:
raise ValueError(f"You should implement the conditonal image of this controlnet: {controlnet}")
assert len(input_images) > 0
kwargs_init_pipeline, kwargs_load_engine, args_run_demo = process_pipeline_args(args)
kwargs_init_pipeline, kwargs_load_engine, args_run_demo = dd_argparse.process_pipeline_args(args)
# Initialize demo
demo = StableDiffusionPipeline(
pipeline_type=PIPELINE_TYPE.CONTROLNET,
demo = pipeline_module.StableDiffusionPipeline(
pipeline_type=(
pipeline_module.PIPELINE_TYPE.CONTROLNET
if args.version != "xl-1.0"
else pipeline_module.PIPELINE_TYPE.XL_CONTROLNET
),
controlnets=args.controlnet_type,
**kwargs_init_pipeline)
**kwargs_init_pipeline,
)
# Load TensorRT engines and pytorch modules
demo.loadEngines(
demo/Diffusion/utils_sd3/__init__.py → demo/Diffusion/demo_diffusion/__init__.py
View File
demo/Diffusion/demo_diffusion/dd_argparse.py
+390
View File
@@ -0,0 +1,390 @@
#
# SPDX-FileCopyrightText: Copyright (c) 1993-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
from __future__ import annotations
import argparse
from typing import Any, Dict, Tuple
import torch
# Define valid optimization levels for TensorRT engine build
VALID_OPTIMIZATION_LEVELS = list(range(6))
def parse_key_value_pairs(string: str) -> Dict[str, str]:
"""Parse a string of comma-separated key-value pairs into a dictionary.
Args:
string (str): A string of comma-separated key-value pairs.
Returns:
Dict[str, str]: Parsed dictionary of key-value pairs.
Example:
>>> parse_key_value_pairs("key1:value1,key2:value2")
{"key1": "value1", "key2": "value2"}
"""
parsed = {}
for key_value_pair in string.split(","):
if not key_value_pair:
continue
key_value_pair = key_value_pair.split(":")
if len(key_value_pair) != 2:
raise argparse.ArgumentTypeError(f"Invalid key-value pair: {key_value_pair}. Must have length 2.")
key, value = key_value_pair
parsed[key] = value
return parsed
def add_arguments(parser):
# Stable Diffusion configuration
parser.add_argument(
"--version",
type=str,
default="1.5",
choices=(
"1.4",
"1.5",
"dreamshaper-7",
"2.0-base",
"2.0",
"2.1-base",
"2.1",
"xl-1.0",
"xl-turbo",
"svd-xt-1.1",
"sd3",
"cascade",
"flux.1-dev",
"flux.1-schnell",
"flux.1-dev-canny",
"flux.1-dev-depth",
),
help="Version of Stable Diffusion",
)
parser.add_argument("prompt", nargs="*", help="Text prompt(s) to guide image generation")
parser.add_argument(
"--negative-prompt", nargs="*", default=[""], help="The negative prompt(s) to guide the image generation."
)
parser.add_argument("--batch-size", type=int, default=1, choices=[1, 2, 4], help="Batch size (repeat prompt)")
parser.add_argument(
"--batch-count", type=int, default=1, help="Number of images to generate in sequence, one at a time."
)
parser.add_argument("--height", type=int, default=512, help="Height of image to generate (must be multiple of 8)")
parser.add_argument("--width", type=int, default=512, help="Height of image to generate (must be multiple of 8)")
parser.add_argument("--denoising-steps", type=int, default=30, help="Number of denoising steps")
parser.add_argument(
"--scheduler",
type=str,
default=None,
choices=("DDIM", "DDPM", "EulerA", "Euler", "LCM", "LMSD", "PNDM", "UniPC", "DDPMWuerstchen", "FlowMatchEuler"),
help="Scheduler for diffusion process",
)
parser.add_argument(
"--guidance-scale",
type=float,
default=7.5,
help="Value of classifier-free guidance scale (must be greater than 1)",
)
parser.add_argument(
"--lora-scale",
type=float,
default=1.0,
help="Controls how much to influence the outputs with the LoRA parameters. (must between 0 and 1)",
)
parser.add_argument(
"--lora-weight",
type=float,
nargs="+",
default=None,
help="The LoRA adapter(s) weights to use with the UNet. (must between 0 and 1)",
)
parser.add_argument(
"--lora-path",
type=str,
nargs="+",
default=None,
help="Path to LoRA adaptor. Ex: 'latent-consistency/lcm-lora-sdv1-5'",
)
# ONNX export
parser.add_argument(
"--onnx-opset",
type=int,
default=19,
choices=range(7, 20),
help="Select ONNX opset version to target for exported models",
)
parser.add_argument("--onnx-dir", default="onnx", help="Output directory for ONNX export")
parser.add_argument(
"--custom-onnx-paths",
type=parse_key_value_pairs,
help=(
"[FLUX only] Custom override paths to pre-exported ONNX model files. These ONNX models are directly used to "
"build TRT engines without further optimization on the ONNX graphs. Paths should be a comma-separated list "
"of <model_name>:<path> pairs. For example: "
"--custom-onnx-paths=transformer:/path/to/transformer.onnx,vae:/path/to/vae.onnx. Call "
"<PipelineClass>.get_model_names(...) for the list of supported model names."
),
)
parser.add_argument(
"--download-onnx-models",
action="store_true",
help=("[FLUX only] Download pre-exported ONNX models"),
)
# Framework model ckpt
parser.add_argument("--framework-model-dir", default="pytorch_model", help="Directory for HF saved models")
# TensorRT engine build
parser.add_argument("--engine-dir", default="engine", help="Output directory for TensorRT engines")
parser.add_argument(
"--custom-engine-paths",
type=parse_key_value_pairs,
help=(
"[FLUX only] Custom override paths to pre-built engine files. Paths should be a comma-separated list of "
"<model_name>:<path> pairs. For example: "
"--custom-onnx-paths=transformer:/path/to/transformer.plan,vae:/path/to/vae.plan. Call "
"<PipelineClass>.get_model_names(...) for the list of supported model names."
),
)
# Quantization configuration.
parser.add_argument("--int8", action="store_true", help="Apply int8 quantization.")
parser.add_argument("--fp8", action="store_true", help="Apply fp8 quantization.")
parser.add_argument("--fp4", action="store_true", help="Apply fp4 quantization.")
parser.add_argument(
"--quantization-level",
type=float,
default=0.0,
choices=[0.0, 1.0, 2.0, 2.5, 3.0, 4.0],
help="int8/fp8 quantization level, 1: CNN, 2: CNN + FFN, 2.5: CNN + FFN + QKV, 3: CNN + Almost all Linear (Including FFN, QKV, Proj and others), 4: CNN + Almost all Linear + fMHA, 0: Default to 2.5 for int8 and 4.0 for fp8.",
)
parser.add_argument(
"--quantization-percentile",
type=float,
default=1.0,
help="Control quantization scaling factors (amax) collecting range, where the minimum amax in range(n_steps * percentile) will be collected. Recommendation: 1.0.",
)
parser.add_argument(
"--quantization-alpha",
type=float,
default=0.8,
help="The alpha parameter for SmoothQuant quantization used for linear layers. Recommendation: 0.8 for SDXL.",
)
parser.add_argument(
"--calibration-size",
type=int,
default=32,
help="The number of steps to use for calibrating the model for quantization. Recommendation: 32, 64, 128 for SDXL",
)
parser.add_argument(
"--optimization-level",
type=int,
default=None,
help=f"Set the builder optimization level to build the engine with. A higher level allows TensorRT to spend more building time for more optimization options. Must be one of {VALID_OPTIMIZATION_LEVELS}.",
)
parser.add_argument(
"--build-static-batch", action="store_true", help="Build TensorRT engines with fixed batch size."
)
parser.add_argument(
"--build-dynamic-shape", action="store_true", help="Build TensorRT engines with dynamic image shapes."
)
parser.add_argument(
"--build-enable-refit", action="store_true", help="Enable Refit option in TensorRT engines during build."
)
parser.add_argument(
"--build-all-tactics", action="store_true", help="Build TensorRT engines using all tactic sources."
)
parser.add_argument(
"--timing-cache", default=None, type=str, help="Path to the precached timing measurements to accelerate build."
)
# TensorRT inference
parser.add_argument(
"--num-warmup-runs", type=int, default=5, help="Number of warmup runs before benchmarking performance"
)
parser.add_argument("--use-cuda-graph", action="store_true", help="Enable cuda graph")
parser.add_argument("--nvtx-profile", action="store_true", help="Enable NVTX markers for performance profiling")
parser.add_argument(
"--torch-inference",
default="",
help="Run inference with PyTorch (using specified compilation mode) instead of TensorRT.",
)
parser.add_argument("--seed", type=int, default=None, help="Seed for random generator to get consistent results")
parser.add_argument("--output-dir", default="output", help="Output directory for logs and image artifacts")
parser.add_argument("--hf-token", type=str, help="HuggingFace API access token for downloading model checkpoints")
parser.add_argument("-v", "--verbose", action="store_true", help="Show verbose output")
return parser
def process_pipeline_args(args: argparse.Namespace) -> Tuple[Dict[str, Any], Dict[str, Any], Tuple]:
"""Validate parsed arguments and process argument values.
Some argument values are resolved or overwritten during processing.
Args:
args (argparse.Namespace): Parsed argument. This is modified in-place.
Returns:
Dict[str, Any]: Keyword arguments for initializing a pipeline. This is only used in legacy pipelines that do not
have factory methods `FromArgs` that construct the pipeline directly from the parsed argument.
Dict[str, Any]: Keyword arguments for calling the `.load_engine` method of the pipeline.
Tuple: Arguments for calling the `.run` method of the pipeline.
"""
# GPU device info
device_info = torch.cuda.get_device_properties(0)
sm_version = device_info.major * 10 + device_info.minor
if args.height % 8 != 0 or args.width % 8 != 0:
raise ValueError(
f"Image height and width have to be divisible by 8 but specified as: {args.image_height} and {args.width}."
)
# Handle batch size
max_batch_size = 4
if args.batch_size > max_batch_size:
raise ValueError(f"Batch size {args.batch_size} is larger than allowed {max_batch_size}.")
if args.use_cuda_graph and (not args.build_static_batch or args.build_dynamic_shape):
raise ValueError(
"Using CUDA graph requires static dimensions. Enable `--build-static-batch` and do not specify `--build-dynamic-shape`"
)
# TensorRT builder optimization level
if args.optimization_level is None:
if args.int8 or args.fp8:
args.optimization_level = 4
else:
args.optimization_level = 3
if args.optimization_level not in VALID_OPTIMIZATION_LEVELS:
raise ValueError(
f"Optimization level {args.optimization_level} not valid. Valid values are: {VALID_OPTIMIZATION_LEVELS}"
)
# Quantized pipeline
# int8 support
if args.int8 and not any(args.version.startswith(prefix) for prefix in ["xl", "1.4", "1.5", "2.1"]):
raise ValueError("int8 quantization is only supported for SDXL, SD1.4, SD1.5 and SD2.1 pipelines.")
# fp8 support
if args.fp8 and not any(
args.version.startswith(prefix) for prefix in ("xl", "1.4", "1.5", "2.1", "flux.1-dev", "flux.1-schnell")
):
raise ValueError("fp8 quantization is only supported for SDXL, SD1.4, SD1.5, SD2.1 and FLUX pipelines.")
if args.fp8 and args.int8:
raise ValueError("Cannot apply both int8 and fp8 quantization, please choose only one.")
if args.fp8 and sm_version < 89:
raise ValueError(
f"Cannot apply FP8 quantization for GPU with compute capability {sm_version / 10.0}. Only Ada and Hopper are supported."
)
# TensorRT ModelOpt quantization level
if args.quantization_level == 0.0:
def override_quant_level(level: float, dtype_str: str):
args.quantization_level = level
print(f"The default quantization level has been set to {level} for {dtype_str}.")
if args.fp8:
override_quant_level(3.0 if args.version in ("1.4", "1.5") else 4.0, "FP8")
# L4 fp8 fMHA on Hopper not yet enabled.
if sm_version == 90 and args.version.startswith("flux"):
override_quant_level(3.0, "FP8")
elif args.int8:
override_quant_level(3.0, "INT8")
if args.quantization_level == 3.0 and args.download_onnx_models:
raise ValueError(
"Transformer ONNX model for Quantization level 3 is not available for download. Please export the quantized Transformer model natively with the removal of --download-onnx-models."
)
if args.fp4:
# FP4 precision is only supported for Flux Pipelines
assert args.version.startswith("flux"), "FP4 precision is only supported for Flux pipelines"
# Handle LoRA
if args.lora_path and not any(args.version.startswith(prefix) for prefix in ("1.5", "2.1", "xl")):
raise ValueError("LoRA adapter support is only supported for SD1.5, SD2.1 and SDXL pipelines")
if args.lora_weight:
for weight in (weight for weight in args.lora_weight if not 0 <= weight <= 1):
raise ValueError(f"LoRA adapter weights must be between 0 and 1, provided {weight}")
if not 0 <= args.lora_scale <= 1:
raise ValueError(f"LoRA scale value must be between 0 and 1, provided {args.lora_scale}")
# Pack arguments
kwargs_init_pipeline = {
"version": args.version,
"max_batch_size": max_batch_size,
"denoising_steps": args.denoising_steps,
"scheduler": args.scheduler,
"guidance_scale": args.guidance_scale,
"output_dir": args.output_dir,
"hf_token": args.hf_token,
"verbose": args.verbose,
"nvtx_profile": args.nvtx_profile,
"use_cuda_graph": args.use_cuda_graph,
"lora_scale": args.lora_scale,
"lora_weight": args.lora_weight,
"lora_path": args.lora_path,
"framework_model_dir": args.framework_model_dir,
"torch_inference": args.torch_inference,
}
kwargs_load_engine = {
"onnx_opset": args.onnx_opset,
"opt_batch_size": args.batch_size,
"opt_image_height": args.height,
"opt_image_width": args.width,
"optimization_level": args.optimization_level,
"static_batch": args.build_static_batch,
"static_shape": not args.build_dynamic_shape,
"enable_all_tactics": args.build_all_tactics,
"enable_refit": args.build_enable_refit,
"timing_cache": args.timing_cache,
"int8": args.int8,
"fp8": args.fp8,
"fp4": args.fp4,
"quantization_level": args.quantization_level,
"quantization_percentile": args.quantization_percentile,
"quantization_alpha": args.quantization_alpha,
"calibration_size": args.calibration_size,
"download_onnx_models": args.download_onnx_models,
}
args_run_demo = (
args.prompt,
args.negative_prompt,
args.height,
args.width,
args.batch_size,
args.batch_count,
args.num_warmup_runs,
args.use_cuda_graph,
)
return kwargs_init_pipeline, kwargs_load_engine, args_run_demo
demo/Diffusion/demo_diffusion/dynamic_import.py
Executable
+30
View File
@@ -0,0 +1,30 @@
#
# SPDX-FileCopyrightText: Copyright (c) 1993-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
import warnings
from importlib import import_module
def import_from_diffusers(model_name, module_name):
try:
module = import_module(module_name)
return getattr(module, model_name)
except ImportError:
warnings.warn(f"Failed to import {module_name}. The {model_name} model will not be available.", ImportWarning)
except AttributeError:
warnings.warn(f"The {model_name} model is not available in the installed version of diffusers.", ImportWarning)
return None
demo/Diffusion/demo_diffusion/engine.py
+337
View File
@@ -0,0 +1,337 @@
#
# SPDX-FileCopyrightText: Copyright (c) 1993-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
import gc
import os
import subprocess
import warnings
from collections import OrderedDict, defaultdict
import numpy as np
import onnx
import tensorrt as trt
import torch
from cuda import cudart
from onnx import numpy_helper
from polygraphy.backend.common import bytes_from_path
from polygraphy.backend.trt import (
engine_from_bytes,
)
TRT_LOGGER = trt.Logger(trt.Logger.ERROR)
# Map of TensorRT dtype -> torch dtype
trt_to_torch_dtype_dict = {
trt.DataType.BOOL: torch.bool,
trt.DataType.UINT8: torch.uint8,
trt.DataType.INT8: torch.int8,
trt.DataType.INT32: torch.int32,
trt.DataType.INT64: torch.int64,
trt.DataType.HALF: torch.float16,
trt.DataType.FLOAT: torch.float32,
trt.DataType.BF16: torch.bfloat16,
}
def _CUASSERT(cuda_ret):
err = cuda_ret[0]
if err != cudart.cudaError_t.cudaSuccess:
raise RuntimeError(
f"CUDA ERROR: {err}, error code reference: https://nvidia.github.io/cuda-python/module/cudart.html#cuda.cudart.cudaError_t"
)
if len(cuda_ret) > 1:
return cuda_ret[1]
return None
def get_refit_weights(state_dict, onnx_opt_path, weight_name_mapping, weight_shape_mapping):
onnx_opt_dir = os.path.dirname(onnx_opt_path)
onnx_opt_model = onnx.load(onnx_opt_path)
# Create initializer data hashes
initializer_hash_mapping = {}
for initializer in onnx_opt_model.graph.initializer:
initializer_data = numpy_helper.to_array(initializer, base_dir=onnx_opt_dir).astype(np.float16)
initializer_hash = hash(initializer_data.data.tobytes())
initializer_hash_mapping[initializer.name] = initializer_hash
refit_weights = OrderedDict()
updated_weight_names = set() # save names of updated weights to refit only the required weights
for wt_name, wt in state_dict.items():
# query initializer to compare
initializer_name = weight_name_mapping[wt_name]
initializer_hash = initializer_hash_mapping[initializer_name]
# get shape transform info
initializer_shape, is_transpose = weight_shape_mapping[wt_name]
if is_transpose:
wt = torch.transpose(wt, 0, 1)
else:
wt = torch.reshape(wt, initializer_shape)
# include weight if hashes differ
wt_hash = hash(wt.cpu().detach().numpy().astype(np.float16).data.tobytes())
if initializer_hash != wt_hash:
updated_weight_names.add(initializer_name)
# Store all weights as the refitter may require unchanged weights too
# docs: https://docs.nvidia.com/deeplearning/tensorrt/developer-guide/index.html#refitting-engine-c
refit_weights[initializer_name] = wt.contiguous()
return refit_weights, updated_weight_names
class Engine:
def __init__(
self,
engine_path,
):
self.engine_path = engine_path
self.engine = None
self.context = None
self.buffers = OrderedDict()
self.tensors = OrderedDict()
self.cuda_graph_instance = None # cuda graph
def __del__(self):
del self.engine
del self.context
del self.buffers
del self.tensors
def refit(self, refit_weights, updated_weight_names):
# Initialize refitter
refitter = trt.Refitter(self.engine, TRT_LOGGER)
refitted_weights = set()
def refit_single_weight(trt_weight_name):
# get weight from state dict
trt_datatype = refitter.get_weights_prototype(trt_weight_name).dtype
refit_weights[trt_weight_name] = refit_weights[trt_weight_name].to(trt_to_torch_dtype_dict[trt_datatype])
# trt.Weight and trt.TensorLocation
trt_wt_tensor = trt.Weights(
trt_datatype, refit_weights[trt_weight_name].data_ptr(), torch.numel(refit_weights[trt_weight_name])
)
trt_wt_location = (
trt.TensorLocation.DEVICE if refit_weights[trt_weight_name].is_cuda else trt.TensorLocation.HOST
)
# apply refit
refitter.set_named_weights(trt_weight_name, trt_wt_tensor, trt_wt_location)
refitted_weights.add(trt_weight_name)
# iterate through all tensorrt refittable weights
for trt_weight_name in refitter.get_all_weights():
if trt_weight_name not in updated_weight_names:
continue
refit_single_weight(trt_weight_name)
# iterate through missing weights required by tensorrt - addresses the case where lora_scale=0
for trt_weight_name in refitter.get_missing_weights():
refit_single_weight(trt_weight_name)
if not refitter.refit_cuda_engine():
print("Error: failed to refit new weights.")
exit(0)
print(f"[I] Total refitted weights {len(refitted_weights)}.")
def build(
self,
onnx_path,
strongly_typed=False,
fp16=True,
bf16=False,
tf32=False,
int8=False,
fp8=False,
input_profile=None,
enable_refit=False,
enable_all_tactics=False,
timing_cache=None,
update_output_names=None,
native_instancenorm=True,
verbose=False,
weight_streaming=False,
builder_optimization_level=3,
precision_constraints='none',
):
print(f"Building TensorRT engine for {onnx_path}: {self.engine_path}")
# Handle weight streaming case: https://docs.nvidia.com/deeplearning/tensorrt/developer-guide/index.html#streaming-weights.
if weight_streaming:
strongly_typed, fp16, bf16, int8, fp8 = True, False, False, False, False
# Base command
build_command = [f"polygraphy convert {onnx_path} --convert-to trt --output {self.engine_path}"]
# Precision flags
build_args = [
"--fp16" if fp16 else "",
"--bf16" if bf16 else "",
"--tf32" if tf32 else "",
"--fp8" if fp8 else "",
"--int8" if int8 else "",
"--strongly-typed" if strongly_typed else "",
]
# Additional arguments
build_args.extend([
"--weight-streaming" if weight_streaming else "",
"--refittable" if enable_refit else "",
"--tactic-sources" if not enable_all_tactics else "",
"--onnx-flags native_instancenorm" if native_instancenorm else "",
f"--builder-optimization-level {builder_optimization_level}",
f"--precision-constraints {precision_constraints}",
])
# Timing cache
if timing_cache:
build_args.extend([
f"--load-timing-cache {timing_cache}",
f"--save-timing-cache {timing_cache}"
])
# Verbosity setting
verbosity = "extra_verbose" if verbose else "error"
build_args.append(f"--verbosity {verbosity}")
# Output names
if update_output_names:
print(f"Updating network outputs to {update_output_names}")
build_args.append(f"--trt-outputs {' '.join(update_output_names)}")
# Input profiles
if input_profile:
profile_args = defaultdict(str)
for name, dims in input_profile.items():
assert len(dims) == 3
profile_args["--trt-min-shapes"] += f"{name}:{str(list(dims[0])).replace(' ', '')} "
profile_args["--trt-opt-shapes"] += f"{name}:{str(list(dims[1])).replace(' ', '')} "
profile_args["--trt-max-shapes"] += f"{name}:{str(list(dims[2])).replace(' ', '')} "
build_args.extend(f"{k} {v}" for k, v in profile_args.items())
# Filter out empty strings and join command
build_args = [arg for arg in build_args if arg]
final_command = ' '.join(build_command + build_args)
# Execute command with improved error handling
try:
print(f"Engine build command: {final_command}")
subprocess.run(final_command, check=True, shell=True)
except subprocess.CalledProcessError as exc:
error_msg = (
f"Failed to build TensorRT engine. Error details:\n"
f"Command: {exc.cmd}\n"
)
raise RuntimeError(error_msg) from exc
def load(self, weight_streaming=False, weight_streaming_budget_percentage=None):
if self.engine is not None:
print(f"[W]: Engine {self.engine_path} already loaded, skip reloading")
return
if not hasattr(self, "engine_bytes_cpu") or self.engine_bytes_cpu is None:
# keep a cpu copy of the engine to reduce reloading time.
print(f"Loading TensorRT engine to cpu bytes: {self.engine_path}")
self.engine_bytes_cpu = bytes_from_path(self.engine_path)
print(f"Loading TensorRT engine from bytes: {self.engine_path}")
self.engine = engine_from_bytes(self.engine_bytes_cpu)
if weight_streaming:
if weight_streaming_budget_percentage is None:
warnings.warn(
f"Weight streaming budget is not set for {self.engine_path}. Weights will not be streamed."
)
else:
self.engine.weight_streaming_budget_v2 = int(
weight_streaming_budget_percentage / 100 * self.engine.streamable_weights_size
)
def unload(self):
if self.engine is not None:
print(f"Unloading TensorRT engine: {self.engine_path}")
del self.engine
self.engine = None
gc.collect()
else:
print(f"[W]: Unload an unloaded engine {self.engine_path}, skip unloading")
def activate(self, device_memory=None):
if device_memory:
self.context = self.engine.create_execution_context_without_device_memory()
self.context.device_memory = device_memory
else:
self.context = self.engine.create_execution_context()
def reactivate(self, device_memory):
assert self.context
self.context.device_memory = device_memory
def deactivate(self):
del self.context
self.context = None
def allocate_buffers(self, shape_dict=None, device="cuda"):
for binding in range(self.engine.num_io_tensors):
name = self.engine.get_tensor_name(binding)
if shape_dict and name in shape_dict:
shape = shape_dict[name]
else:
shape = self.engine.get_tensor_shape(name)
print(
f"[W]: {self.engine_path}: Could not find '{name}' in shape dict {shape_dict}. Using shape {shape} inferred from the engine."
)
if self.engine.get_tensor_mode(name) == trt.TensorIOMode.INPUT:
self.context.set_input_shape(name, shape)
dtype = trt_to_torch_dtype_dict[self.engine.get_tensor_dtype(name)]
tensor = torch.empty(tuple(shape), dtype=dtype).to(device=device)
self.tensors[name] = tensor
def deallocate_buffers(self):
for idx in range(self.engine.num_io_tensors):
binding = self.engine[idx]
del self.tensors[binding]
def infer(self, feed_dict, stream, use_cuda_graph=False):
for name, buf in feed_dict.items():
self.tensors[name].copy_(buf)
for name, tensor in self.tensors.items():
self.context.set_tensor_address(name, tensor.data_ptr())
if use_cuda_graph:
if self.cuda_graph_instance is not None:
_CUASSERT(cudart.cudaGraphLaunch(self.cuda_graph_instance, stream))
_CUASSERT(cudart.cudaStreamSynchronize(stream))
else:
# do inference before CUDA graph capture
noerror = self.context.execute_async_v3(stream)
if not noerror:
raise ValueError(f"ERROR: inference of {self.engine_path} failed.")
# capture cuda graph
_CUASSERT(
cudart.cudaStreamBeginCapture(stream, cudart.cudaStreamCaptureMode.cudaStreamCaptureModeGlobal)
)
self.context.execute_async_v3(stream)
self.graph = _CUASSERT(cudart.cudaStreamEndCapture(stream))
self.cuda_graph_instance = _CUASSERT(cudart.cudaGraphInstantiate(self.graph, 0))
else:
noerror = self.context.execute_async_v3(stream)
if not noerror:
raise ValueError(f"ERROR: inference of {self.engine_path} failed.")
return self.tensors
demo/Diffusion/demo_diffusion/image/__init__.py
+34
View File
@@ -0,0 +1,34 @@
#
# SPDX-FileCopyrightText: Copyright (c) 1993-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
from demo_diffusion.image.load import (
download_image,
prepare_mask_and_masked_image,
preprocess_image,
save_image,
)
from demo_diffusion.image.resize import resize_with_antialiasing
from demo_diffusion.image.video import tensor2vid
__all__ = [
"preprocess_image",
"prepare_mask_and_masked_image",
"download_image",
"save_image",
"resize_with_antialiasing",
"tensor2vid",
]
demo/Diffusion/demo_diffusion/image/load.py
+78
View File
@@ -0,0 +1,78 @@
#
# SPDX-FileCopyrightText: Copyright (c) 1993-2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
import os
import random
from io import BytesIO
import numpy as np
import requests
import torch
from PIL import Image
def preprocess_image(image):
"""
image: torch.Tensor
"""
w, h = image.size
w, h = map(lambda x: x - x % 32, (w, h)) # resize to integer multiple of 32
image = image.resize((w, h))
image = np.array(image).astype(np.float32) / 255.0
image = image[None].transpose(0, 3, 1, 2)
image = torch.from_numpy(image).contiguous()
return 2.0 * image - 1.0
def prepare_mask_and_masked_image(image, mask):
"""
image: PIL.Image.Image
mask: PIL.Image.Image
"""
if isinstance(image, Image.Image):
image = np.array(image.convert("RGB"))
image = image[None].transpose(0, 3, 1, 2)
image = torch.from_numpy(image).to(dtype=torch.float32).contiguous() / 127.5 - 1.0
if isinstance(mask, Image.Image):
mask = np.array(mask.convert("L"))
mask = mask.astype(np.float32) / 255.0
mask = mask[None, None]
mask[mask < 0.5] = 0
mask[mask >= 0.5] = 1
mask = torch.from_numpy(mask).to(dtype=torch.float32).contiguous()
masked_image = image * (mask < 0.5)
return mask, masked_image
def download_image(url):
response = requests.get(url)
return Image.open(BytesIO(response.content)).convert("RGB")
def save_image(images, image_path_dir, image_name_prefix, image_name_suffix):
"""
Save the generated images to png files.
"""
for i in range(images.shape[0]):
image_path = os.path.join(
image_path_dir,
f"{image_name_prefix}{i + 1}-{random.randint(1000, 9999)}-{image_name_suffix}.png",
)
print(f"Saving image {i+1} / {images.shape[0]} to: {image_path}")
Image.fromarray(images[i]).save(image_path)
demo/Diffusion/demo_diffusion/image/resize.py
+123
View File
@@ -0,0 +1,123 @@
# Copyright 2024 The HuggingFace Team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import torch
# Taken from https://github.com/huggingface/diffusers/blob/be62c85cd973f2001ab8c5d8919a9a6811fc7e43/src/diffusers/pipelines/stable_video_diffusion/pipeline_stable_video_diffusion.py#L633
def resize_with_antialiasing(input, size, interpolation="bicubic", align_corners=True):
h, w = input.shape[-2:]
factors = (h / size[0], w / size[1])
# First, we have to determine sigma
# Taken from skimage: https://github.com/scikit-image/scikit-image/blob/v0.19.2/skimage/transform/_warps.py#L171
sigmas = (
max((factors[0] - 1.0) / 2.0, 0.001),
max((factors[1] - 1.0) / 2.0, 0.001),
)
# Now kernel size. Good results are for 3 sigma, but that is kind of slow. Pillow uses 1 sigma
# https://github.com/python-pillow/Pillow/blob/master/src/libImaging/Resample.c#L206
# But they do it in the 2 passes, which gives better results. Let's try 2 sigmas for now
ks = int(max(2.0 * 2 * sigmas[0], 3)), int(max(2.0 * 2 * sigmas[1], 3))
# Make sure it is odd
if (ks[0] % 2) == 0:
ks = ks[0] + 1, ks[1]
if (ks[1] % 2) == 0:
ks = ks[0], ks[1] + 1
input = _gaussian_blur2d(input, ks, sigmas)
output = torch.nn.functional.interpolate(input, size=size, mode=interpolation, align_corners=align_corners)
return output
def _compute_padding(kernel_size):
"""Compute padding tuple."""
# 4 or 6 ints: (padding_left, padding_right,padding_top,padding_bottom)
# https://pytorch.org/docs/stable/nn.html#torch.nn.functional.pad
if len(kernel_size) < 2:
raise AssertionError(kernel_size)
computed = [k - 1 for k in kernel_size]
# for even kernels we need to do asymmetric padding :(
out_padding = 2 * len(kernel_size) * [0]
for i in range(len(kernel_size)):
computed_tmp = computed[-(i + 1)]
pad_front = computed_tmp // 2
pad_rear = computed_tmp - pad_front
out_padding[2 * i + 0] = pad_front
out_padding[2 * i + 1] = pad_rear
return out_padding
def _filter2d(input, kernel):
# prepare kernel
b, c, h, w = input.shape
tmp_kernel = kernel[:, None, ...].to(device=input.device, dtype=input.dtype)
tmp_kernel = tmp_kernel.expand(-1, c, -1, -1)
height, width = tmp_kernel.shape[-2:]
padding_shape: list[int] = _compute_padding([height, width])
input = torch.nn.functional.pad(input, padding_shape, mode="reflect")
# kernel and input tensor reshape to align element-wise or batch-wise params
tmp_kernel = tmp_kernel.reshape(-1, 1, height, width)
input = input.view(-1, tmp_kernel.size(0), input.size(-2), input.size(-1))
# convolve the tensor with the kernel.
output = torch.nn.functional.conv2d(input, tmp_kernel, groups=tmp_kernel.size(0), padding=0, stride=1)
out = output.view(b, c, h, w)
return out
def _gaussian(window_size: int, sigma):
if isinstance(sigma, float):
sigma = torch.tensor([[sigma]])
batch_size = sigma.shape[0]
x = (torch.arange(window_size, device=sigma.device, dtype=sigma.dtype) - window_size // 2).expand(batch_size, -1)
if window_size % 2 == 0:
x = x + 0.5
gauss = torch.exp(-x.pow(2.0) / (2 * sigma.pow(2.0)))
return gauss / gauss.sum(-1, keepdim=True)
def _gaussian_blur2d(input, kernel_size, sigma):
if isinstance(sigma, tuple):
sigma = torch.tensor([sigma], dtype=input.dtype)
else:
sigma = sigma.to(dtype=input.dtype)
ky, kx = int(kernel_size[0]), int(kernel_size[1])
bs = sigma.shape[0]
kernel_x = _gaussian(kx, sigma[:, 1].view(bs, 1))
kernel_y = _gaussian(ky, sigma[:, 0].view(bs, 1))
out_x = _filter2d(input, kernel_x[..., None, :])
out = _filter2d(out_x, kernel_y[..., None])
return out
demo/Diffusion/demo_diffusion/image/video.py
+37
View File
@@ -0,0 +1,37 @@
#
# Copyright (c) Alibaba, Inc. and its affiliates.
# SPDX-FileCopyrightText: Copyright (c) 1993-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
import torch
# Not a contribution
# Changes made by NVIDIA CORPORATION & AFFILIATES enabling tensor2vid or otherwise documented as
# NVIDIA-proprietary are not a contribution and subject to the terms and conditions at the top of the file
def tensor2vid(video: torch.Tensor, processor, output_type="np"):
# Based on:
# https://github.com/modelscope/modelscope/blob/1509fdb973e5871f37148a4b5e5964cafd43e64d/modelscope/pipelines/multi_modal/text_to_video_synthesis_pipeline.py#L78
batch_size, channels, num_frames, height, width = video.shape
outputs = []
for batch_idx in range(batch_size):
batch_vid = video[batch_idx].permute(1, 0, 2, 3)
batch_output = processor.postprocess(batch_vid, output_type)
outputs.append(batch_output)
return outputs
demo/Diffusion/demo_diffusion/model/__init__.py
+96
View File
@@ -0,0 +1,96 @@
#
# SPDX-FileCopyrightText: Copyright (c) 1993-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
from demo_diffusion.model.base_model import BaseModel
from demo_diffusion.model.clip import (
CLIPImageProcessorModel,
CLIPModel,
CLIPVisionWithProjModel,
CLIPWithProjModel,
SD3_CLIPGModel,
SD3_CLIPLModel,
SD3_T5XXLModel,
get_clip_embedding_dim,
)
from demo_diffusion.model.diffusion_transformer import (
FluxTransformerModel,
SD3_MMDiTModel,
)
from demo_diffusion.model.gan import VQGANModel
from demo_diffusion.model.load import unload_torch_model
from demo_diffusion.model.lora import LoraLoader, merge_loras
from demo_diffusion.model.scheduler import make_scheduler
from demo_diffusion.model.t5 import T5Model
from demo_diffusion.model.tokenizer import make_tokenizer
from demo_diffusion.model.unet import (
UNetCascadeModel,
UNetModel,
UNetTemporalModel,
UNetXLModel,
UNetXLModelControlNet,
)
from demo_diffusion.model.vae import (
SD3_VAEDecoderModel,
SD3_VAEEncoderModel,
TorchVAEEncoder,
VAEDecTemporalModel,
VAEEncoderModel,
VAEModel,
)
__all__ = [
# base_model
"BaseModel",
# clip
"get_clip_embedding_dim",
"CLIPModel",
"CLIPWithProjModel",
"SD3_CLIPGModel",
"SD3_CLIPLModel",
"SD3_T5XXLModel",
"CLIPVisionWithProjModel",
"CLIPImageProcessorModel",
# diffusion_transformer
"SD3_MMDiTModel",
"FluxTransformerModel",
# gan
"VQGANModel",
# lora
"LoraLoader",
"merge_loras",
# scheduler
"make_scheduler",
# t5
"T5Model",
# tokenizer
"make_tokenizer",
# unet
"UNetModel",
"UNetXLModel",
"UNetXLModelControlNet",
"UNetTemporalModel",
"UNetCascadeModel",
# vae
"VAEModel",
"SD3_VAEDecoderModel",
"VAEDecTemporalModel",
"TorchVAEEncoder",
"VAEEncoderModel",
"SD3_VAEEncoderModel",
# load
"unload_torch_model",
]
demo/Diffusion/demo_diffusion/model/base_model.py
+277
View File
@@ -0,0 +1,277 @@
#
# SPDX-FileCopyrightText: Copyright (c) 1993-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
import json
import os
import numpy as np
import torch
from diffusers import DiffusionPipeline
import onnx
from demo_diffusion.model import load, optimizer
from demo_diffusion.model.lora import merge_loras
from onnx import numpy_helper
class BaseModel:
def __init__(
self,
version="1.5",
pipeline=None,
device="cuda",
hf_token="",
verbose=True,
framework_model_dir="pytorch_model",
fp16=False,
tf32=False,
bf16=False,
int8=False,
fp8=False,
max_batch_size=16,
text_maxlen=77,
embedding_dim=768,
compression_factor=8,
):
self.name = self.__class__.__name__
self.pipeline = pipeline.name
self.version = version
self.path = load.get_path(version, pipeline)
self.device = device
self.hf_token = hf_token
self.hf_safetensor = not (pipeline.is_inpaint() and version in ("1.4", "1.5"))
self.verbose = verbose
self.framework_model_dir = framework_model_dir
self.fp16 = fp16
self.tf32 = tf32
self.bf16 = bf16
self.int8 = int8
self.fp8 = fp8
self.compression_factor = compression_factor
self.min_batch = 1
self.max_batch = max_batch_size
self.min_image_shape = 256 # min image resolution: 256x256
self.max_image_shape = 1344 # max image resolution: 1344x1344
self.min_latent_shape = self.min_image_shape // self.compression_factor
self.max_latent_shape = self.max_image_shape // self.compression_factor
self.text_maxlen = text_maxlen
self.embedding_dim = embedding_dim
self.extra_output_names = []
self.do_constant_folding = True
def get_pipeline(self):
model_opts = {"variant": "fp16", "torch_dtype": torch.float16} if self.fp16 else {}
model_opts = {"torch_dtype": torch.bfloat16} if self.bf16 else model_opts
return DiffusionPipeline.from_pretrained(
self.path,
use_safetensors=self.hf_safetensor,
token=self.hf_token,
**model_opts,
).to(self.device)
def get_model(self, torch_inference=""):
pass
def get_input_names(self):
pass
def get_output_names(self):
pass
def get_dynamic_axes(self):
return None
def get_sample_input(self, batch_size, image_height, image_width, static_shape):
pass
def get_input_profile(self, batch_size, image_height, image_width, static_batch, static_shape):
return None
def get_shape_dict(self, batch_size, image_height, image_width):
return None
# Helper utility for ONNX export
def export_onnx(
self,
onnx_path,
onnx_opt_path,
onnx_opset,
opt_image_height,
opt_image_width,
custom_model=None,
enable_lora_merge=False,
static_shape=False,
lora_loader=None,
):
onnx_opt_graph = None
# Export optimized ONNX model (if missing)
if not os.path.exists(onnx_opt_path):
if not os.path.exists(onnx_path):
print(f"[I] Exporting ONNX model: {onnx_path}")
def export_onnx(model):
if enable_lora_merge:
assert lora_loader is not None
model = merge_loras(model, lora_loader)
inputs = self.get_sample_input(1, opt_image_height, opt_image_width, static_shape)
torch.onnx.export(
model,
inputs,
onnx_path,
export_params=True,
opset_version=onnx_opset,
do_constant_folding=self.do_constant_folding,
input_names=self.get_input_names(),
output_names=self.get_output_names(),
dynamic_axes=self.get_dynamic_axes(),
verbose=False,
)
if custom_model:
with torch.inference_mode():
export_onnx(custom_model)
else:
# WAR: Enable autocast for BF16 Stable Cascade pipeline
do_autocast = True if self.version == "cascade" and self.bf16 else False
with torch.inference_mode(), torch.autocast("cuda", enabled=do_autocast):
export_onnx(self.get_model())
else:
print(f"[I] Found cached ONNX model: {onnx_path}")
print(f"[I] Optimizing ONNX model: {onnx_opt_path}")
onnx_opt_graph = self.optimize(onnx.load(onnx_path))
if load.onnx_graph_needs_external_data(onnx_opt_graph):
onnx.save_model(
onnx_opt_graph,
onnx_opt_path,
save_as_external_data=True,
all_tensors_to_one_file=True,
convert_attribute=False,
)
else:
onnx.save(onnx_opt_graph, onnx_opt_path)
else:
print(f"[I] Found cached optimized ONNX model: {onnx_opt_path} ")
# Helper utility for weights map
def export_weights_map(self, onnx_opt_path, weights_map_path):
if not os.path.exists(weights_map_path):
onnx_opt_dir = os.path.dirname(onnx_opt_path)
onnx_opt_model = onnx.load(onnx_opt_path)
state_dict = self.get_model().state_dict()
# Create initializer data hashes
initializer_hash_mapping = {}
for initializer in onnx_opt_model.graph.initializer:
initializer_data = numpy_helper.to_array(initializer, base_dir=onnx_opt_dir).astype(np.float16)
initializer_hash = hash(initializer_data.data.tobytes())
initializer_hash_mapping[initializer.name] = (initializer_hash, initializer_data.shape)
weights_name_mapping = {}
weights_shape_mapping = {}
# set to keep track of initializers already added to the name_mapping dict
initializers_mapped = set()
for wt_name, wt in state_dict.items():
# get weight hash
wt = wt.cpu().detach().numpy().astype(np.float16)
wt_hash = hash(wt.data.tobytes())
wt_t_hash = hash(np.transpose(wt).data.tobytes())
for initializer_name, (initializer_hash, initializer_shape) in initializer_hash_mapping.items():
# Due to constant folding, some weights are transposed during export
# To account for the transpose op, we compare the initializer hash to the
# hash for the weight and its transpose
if wt_hash == initializer_hash or wt_t_hash == initializer_hash:
# The assert below ensures there is a 1:1 mapping between
# PyTorch and ONNX weight names. It can be removed in cases where 1:many
# mapping is found and name_mapping[wt_name] = list()
assert initializer_name not in initializers_mapped
weights_name_mapping[wt_name] = initializer_name
initializers_mapped.add(initializer_name)
is_transpose = False if wt_hash == initializer_hash else True
weights_shape_mapping[wt_name] = (initializer_shape, is_transpose)
# Sanity check: Were any weights not matched
if wt_name not in weights_name_mapping:
print(f"[I] PyTorch weight {wt_name} not matched with any ONNX initializer")
print(f"[I] {len(weights_name_mapping.keys())} PyTorch weights were matched with ONNX initializers")
assert weights_name_mapping.keys() == weights_shape_mapping.keys()
with open(weights_map_path, "w") as fp:
json.dump([weights_name_mapping, weights_shape_mapping], fp)
else:
print(f"[I] Found cached weights map: {weights_map_path} ")
def optimize(self, onnx_graph, return_onnx=True, **kwargs):
opt = optimizer.Optimizer(onnx_graph, verbose=self.verbose)
opt.info(self.name + ": original")
opt.cleanup()
opt.info(self.name + ": cleanup")
if kwargs.get("modify_fp8_graph", False):
is_fp16_io = kwargs.get("is_fp16_io", True)
opt.modify_fp8_graph(is_fp16_io=is_fp16_io)
opt.info(self.name + ": modify fp8 graph")
if self.version.startswith("flux.1") and self.fp8:
opt.flux_convert_rope_weight_type()
opt.info(self.name + ": convert rope weight type for fp8 flux")
opt.fold_constants()
opt.info(self.name + ": fold constants")
opt.infer_shapes()
opt.info(self.name + ": shape inference")
if kwargs.get("fuse_mha_qkv_int8", False):
opt.fuse_mha_qkv_int8_sq()
opt.info(self.name + ": fuse QKV nodes")
onnx_opt_graph = opt.cleanup(return_onnx=return_onnx)
opt.info(self.name + ": finished")
return onnx_opt_graph
def check_dims(self, batch_size, image_height, image_width):
assert batch_size >= self.min_batch and batch_size <= self.max_batch
latent_height = image_height // self.compression_factor
latent_width = image_width // self.compression_factor
assert latent_height >= self.min_latent_shape and latent_height <= self.max_latent_shape
assert latent_width >= self.min_latent_shape and latent_width <= self.max_latent_shape
return (latent_height, latent_width)
def get_minmax_dims(self, batch_size, image_height, image_width, static_batch, static_shape):
min_batch = batch_size if static_batch else self.min_batch
max_batch = batch_size if static_batch else self.max_batch
latent_height = image_height // self.compression_factor
latent_width = image_width // self.compression_factor
min_image_height = image_height if static_shape else self.min_image_shape
max_image_height = image_height if static_shape else self.max_image_shape
min_image_width = image_width if static_shape else self.min_image_shape
max_image_width = image_width if static_shape else self.max_image_shape
min_latent_height = latent_height if static_shape else self.min_latent_shape
max_latent_height = latent_height if static_shape else self.max_latent_shape
min_latent_width = latent_width if static_shape else self.min_latent_shape
max_latent_width = latent_width if static_shape else self.max_latent_shape
return (
min_batch,
max_batch,
min_image_height,
max_image_height,
min_image_width,
max_image_width,
min_latent_height,
max_latent_height,
min_latent_width,
max_latent_width,
)
demo/Diffusion/demo_diffusion/model/clip.py
+549
View File
@@ -0,0 +1,549 @@
#
# SPDX-FileCopyrightText: Copyright (c) 1993-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
import os
import torch
from huggingface_hub import hf_hub_download
from safetensors import safe_open
from transformers import (
CLIPImageProcessor,
CLIPTextModel,
CLIPTextModelWithProjection,
CLIPVisionModelWithProjection,
)
from demo_diffusion.model import base_model, load, optimizer
from demo_diffusion.utils_sd3.other_impls import (
SDClipModel,
SDXLClipG,
T5XXLModel,
load_into,
)
def get_clipwithproj_embedding_dim(version: str, pipeline: str) -> int:
"""Return the embedding dimension of a CLIP with projection model."""
if version in ("xl-1.0", "xl-turbo", "cascade"):
return 1280
else:
raise ValueError(f"Invalid version {version} + pipeline {pipeline}")
def get_clip_embedding_dim(version, pipeline):
if version in (
"1.4",
"1.5",
"dreamshaper-7",
"flux.1-dev",
"flux.1-schnell",
"flux.1-dev-canny",
"flux.1-dev-depth",
):
return 768
elif version in ("2.0", "2.0-base", "2.1", "2.1-base"):
return 1024
elif version in ("xl-1.0", "xl-turbo") and pipeline.is_sd_xl_base():
return 768
elif version in ("sd3"):
return 4096
else:
raise ValueError(f"Invalid version {version} + pipeline {pipeline}")
class CLIPModel(base_model.BaseModel):
def __init__(
self,
version,
pipeline,
device,
hf_token,
verbose,
framework_model_dir,
max_batch_size,
embedding_dim,
fp16=False,
tf32=False,
bf16=False,
output_hidden_states=False,
keep_pooled_output=False,
subfolder="text_encoder",
):
super(CLIPModel, self).__init__(
version,
pipeline,
device=device,
hf_token=hf_token,
verbose=verbose,
framework_model_dir=framework_model_dir,
fp16=fp16,
tf32=tf32,
bf16=bf16,
max_batch_size=max_batch_size,
embedding_dim=embedding_dim,
)
self.subfolder = subfolder
self.hidden_layer_offset = 0 if pipeline.is_cascade() else -1
self.keep_pooled_output = keep_pooled_output
# Output the final hidden state
if output_hidden_states:
self.extra_output_names = ["hidden_states"]
def get_model(self, torch_inference=""):
model_opts = (
{"torch_dtype": torch.float16} if self.fp16 else {"torch_dtype": torch.bfloat16} if self.bf16 else {}
)
clip_model_dir = load.get_checkpoint_dir(self.framework_model_dir, self.version, self.pipeline, self.subfolder)
if not load.is_model_cached(clip_model_dir, model_opts, self.hf_safetensor, model_name="model"):
model = CLIPTextModel.from_pretrained(
self.path,
subfolder=self.subfolder,
use_safetensors=self.hf_safetensor,
token=self.hf_token,
**model_opts,
).to(self.device)
model.save_pretrained(clip_model_dir, **model_opts)
else:
print(f"[I] Load CLIPTextModel model from: {clip_model_dir}")
model = CLIPTextModel.from_pretrained(clip_model_dir, **model_opts).to(self.device)
model = optimizer.optimize_checkpoint(model, torch_inference)
return model
def get_input_names(self):
return ["input_ids"]
def get_output_names(self):
output_names = ["text_embeddings"]
if self.keep_pooled_output:
output_names += ["pooled_embeddings"]
return output_names
def get_dynamic_axes(self):
dynamic_axes = {
"input_ids": {0: "B"},
"text_embeddings": {0: "B"},
}
if self.keep_pooled_output:
dynamic_axes["pooled_embeddings"] = {0: "B"}
return dynamic_axes
def get_input_profile(self, batch_size, image_height, image_width, static_batch, static_shape):
self.check_dims(batch_size, image_height, image_width)
min_batch, max_batch, _, _, _, _, _, _, _, _ = self.get_minmax_dims(
batch_size, image_height, image_width, static_batch, static_shape
)
return {
"input_ids": [(min_batch, self.text_maxlen), (batch_size, self.text_maxlen), (max_batch, self.text_maxlen)]
}
def get_shape_dict(self, batch_size, image_height, image_width):
self.check_dims(batch_size, image_height, image_width)
output = {
"input_ids": (batch_size, self.text_maxlen),
"text_embeddings": (batch_size, self.text_maxlen, self.embedding_dim),
}
if self.keep_pooled_output:
output["pooled_embeddings"] = (batch_size, self.embedding_dim)
if "hidden_states" in self.extra_output_names:
output["hidden_states"] = (batch_size, self.text_maxlen, self.embedding_dim)
return output
def get_sample_input(self, batch_size, image_height, image_width, static_shape):
self.check_dims(batch_size, image_height, image_width)
return torch.zeros(batch_size, self.text_maxlen, dtype=torch.int32, device=self.device)
def optimize(self, onnx_graph):
opt = optimizer.Optimizer(onnx_graph, verbose=self.verbose)
opt.info(self.name + ": original")
keep_outputs = [0, 1] if self.keep_pooled_output else [0]
opt.select_outputs(keep_outputs)
opt.cleanup()
opt.fold_constants()
opt.info(self.name + ": fold constants")
opt.infer_shapes()
opt.info(self.name + ": shape inference")
opt.select_outputs(keep_outputs, names=self.get_output_names()) # rename network outputs
opt.info(self.name + ": rename network output(s)")
opt_onnx_graph = opt.cleanup(return_onnx=True)
if "hidden_states" in self.extra_output_names:
opt_onnx_graph = opt.clip_add_hidden_states(self.hidden_layer_offset, return_onnx=True)
opt.info(self.name + ": added hidden_states")
opt.info(self.name + ": finished")
return opt_onnx_graph
class CLIPWithProjModel(CLIPModel):
def __init__(
self,
version,
pipeline,
device,
hf_token,
verbose,
framework_model_dir,
fp16=False,
bf16=False,
max_batch_size=16,
output_hidden_states=False,
subfolder="text_encoder_2",
):
super(CLIPWithProjModel, self).__init__(
version,
pipeline,
device=device,
hf_token=hf_token,
verbose=verbose,
framework_model_dir=framework_model_dir,
fp16=fp16,
bf16=bf16,
max_batch_size=max_batch_size,
embedding_dim=get_clipwithproj_embedding_dim(version, pipeline),
output_hidden_states=output_hidden_states,
)
self.subfolder = subfolder
def get_model(self, torch_inference=""):
model_opts = {"variant": "bf16", "torch_dtype": torch.bfloat16} if self.bf16 else {}
clip_model_dir = load.get_checkpoint_dir(self.framework_model_dir, self.version, self.pipeline, self.subfolder)
if not load.is_model_cached(clip_model_dir, model_opts, self.hf_safetensor, model_name="model"):
model = CLIPTextModelWithProjection.from_pretrained(
self.path,
subfolder=self.subfolder,
use_safetensors=self.hf_safetensor,
token=self.hf_token,
**model_opts,
).to(self.device)
model.save_pretrained(clip_model_dir, **model_opts)
else:
print(f"[I] Load CLIPTextModelWithProjection model from: {clip_model_dir}")
model = CLIPTextModelWithProjection.from_pretrained(clip_model_dir, **model_opts).to(self.device)
model = optimizer.optimize_checkpoint(model, torch_inference)
return model
def get_input_names(self):
return ["input_ids", "attention_mask"]
def get_output_names(self):
return ["text_embeddings"]
def get_dynamic_axes(self):
return {"input_ids": {0: "B"}, "attention_mask": {0: "B"}, "text_embeddings": {0: "B"}}
def get_input_profile(self, batch_size, image_height, image_width, static_batch, static_shape):
self.check_dims(batch_size, image_height, image_width)
min_batch, max_batch, _, _, _, _, _, _, _, _ = self.get_minmax_dims(
batch_size, image_height, image_width, static_batch, static_shape
)
return {
"input_ids": [(min_batch, self.text_maxlen), (batch_size, self.text_maxlen), (max_batch, self.text_maxlen)],
"attention_mask": [
(min_batch, self.text_maxlen),
(batch_size, self.text_maxlen),
(max_batch, self.text_maxlen),
],
}
def get_shape_dict(self, batch_size, image_height, image_width):
self.check_dims(batch_size, image_height, image_width)
output = {
"input_ids": (batch_size, self.text_maxlen),
"attention_mask": (batch_size, self.text_maxlen),
"text_embeddings": (batch_size, self.embedding_dim),
}
if "hidden_states" in self.extra_output_names:
output["hidden_states"] = (batch_size, self.text_maxlen, self.embedding_dim)
return output
def get_sample_input(self, batch_size, image_height, image_width, static_shape):
self.check_dims(batch_size, image_height, image_width)
return (
torch.zeros(batch_size, self.text_maxlen, dtype=torch.int32, device=self.device),
torch.zeros(batch_size, self.text_maxlen, dtype=torch.int32, device=self.device),
)
class SD3_CLIPGModel(CLIPModel):
def __init__(
self,
version,
pipeline,
device,
hf_token,
verbose,
framework_model_dir,
max_batch_size,
embedding_dim=None,
fp16=False,
pooled_output=False,
):
self.CLIPG_CONFIG = {
"hidden_act": "gelu",
"hidden_size": 1280,
"intermediate_size": 5120,
"num_attention_heads": 20,
"num_hidden_layers": 32,
}
super(SD3_CLIPGModel, self).__init__(
version,
pipeline,
device=device,
hf_token=hf_token,
verbose=verbose,
framework_model_dir=framework_model_dir,
fp16=fp16,
max_batch_size=max_batch_size,
embedding_dim=self.CLIPG_CONFIG["hidden_size"] if embedding_dim is None else embedding_dim,
)
self.subfolder = "text_encoders"
if pooled_output:
self.extra_output_names = ["pooled_output"]
def get_model(self, torch_inference=""):
clip_g_model_dir = load.get_checkpoint_dir(
self.framework_model_dir, self.version, self.pipeline, self.subfolder
)
clip_g_filename = "clip_g.safetensors"
clip_g_model_path = f"{clip_g_model_dir}/{clip_g_filename}"
if not os.path.exists(clip_g_model_path):
hf_hub_download(
repo_id=self.path,
filename=clip_g_filename,
local_dir=load.get_checkpoint_dir(self.framework_model_dir, self.version, self.pipeline, ""),
subfolder=self.subfolder,
)
with safe_open(clip_g_model_path, framework="pt", device=self.device) as f:
dtype = torch.float16 if self.fp16 else torch.float32
model = SDXLClipG(self.CLIPG_CONFIG, device=self.device, dtype=dtype)
load_into(f, model.transformer, "", self.device, dtype)
model = optimizer.optimize_checkpoint(model, torch_inference)
return model
def get_shape_dict(self, batch_size, image_height, image_width):
self.check_dims(batch_size, image_height, image_width)
output = {
"input_ids": (batch_size, self.text_maxlen),
"text_embeddings": (batch_size, self.text_maxlen, self.embedding_dim),
}
if "pooled_output" in self.extra_output_names:
output["pooled_output"] = (batch_size, self.embedding_dim)
return output
def optimize(self, onnx_graph):
opt = optimizer.Optimizer(onnx_graph, verbose=self.verbose)
opt.info(self.name + ": original")
opt.select_outputs([0, 1])
opt.cleanup()
opt.fold_constants()
opt.info(self.name + ": fold constants")
opt.infer_shapes()
opt.info(self.name + ": shape inference")
opt.select_outputs([0, 1], names=["text_embeddings", "pooled_output"]) # rename network output
opt.info(self.name + ": rename output[0] and output[1]")
opt_onnx_graph = opt.cleanup(return_onnx=True)
opt.info(self.name + ": finished")
return opt_onnx_graph
class SD3_CLIPLModel(SD3_CLIPGModel):
def __init__(
self,
version,
pipeline,
device,
hf_token,
verbose,
framework_model_dir,
max_batch_size,
fp16=False,
pooled_output=False,
):
self.CLIPL_CONFIG = {
"hidden_act": "quick_gelu",
"hidden_size": 768,
"intermediate_size": 3072,
"num_attention_heads": 12,
"num_hidden_layers": 12,
}
super(SD3_CLIPLModel, self).__init__(
version,
pipeline,
device=device,
hf_token=hf_token,
verbose=verbose,
framework_model_dir=framework_model_dir,
fp16=fp16,
max_batch_size=max_batch_size,
embedding_dim=self.CLIPL_CONFIG["hidden_size"],
)
self.subfolder = "text_encoders"
if pooled_output:
self.extra_output_names = ["pooled_output"]
def get_model(self, torch_inference=""):
clip_l_model_dir = load.get_checkpoint_dir(
self.framework_model_dir, self.version, self.pipeline, self.subfolder
)
clip_l_filename = "clip_l.safetensors"
clip_l_model_path = f"{clip_l_model_dir}/{clip_l_filename}"
if not os.path.exists(clip_l_model_path):
hf_hub_download(
repo_id=self.path,
filename=clip_l_filename,
local_dir=load.get_checkpoint_dir(self.framework_model_dir, self.version, self.pipeline, ""),
subfolder=self.subfolder,
)
with safe_open(clip_l_model_path, framework="pt", device=self.device) as f:
dtype = torch.float16 if self.fp16 else torch.float32
model = SDClipModel(
layer="hidden",
layer_idx=-2,
device=self.device,
dtype=dtype,
layer_norm_hidden_state=False,
return_projected_pooled=False,
textmodel_json_config=self.CLIPL_CONFIG,
)
load_into(f, model.transformer, "", self.device, dtype)
model = optimizer.optimize_checkpoint(model, torch_inference)
return model
# NOTE: For legacy reasons, even though this is a T5 model, it inherits from CLIPModel.
class SD3_T5XXLModel(CLIPModel):
def __init__(
self,
version,
pipeline,
device,
hf_token,
verbose,
framework_model_dir,
max_batch_size,
embedding_dim,
fp16=False,
):
super(SD3_T5XXLModel, self).__init__(
version,
pipeline,
device=device,
hf_token=hf_token,
verbose=verbose,
framework_model_dir=framework_model_dir,
fp16=fp16,
max_batch_size=max_batch_size,
embedding_dim=embedding_dim,
)
self.T5_CONFIG = {"d_ff": 10240, "d_model": 4096, "num_heads": 64, "num_layers": 24, "vocab_size": 32128}
self.subfolder = "text_encoders"
def get_model(self, torch_inference=""):
t5xxl_model_dir = load.get_checkpoint_dir(self.framework_model_dir, self.version, self.pipeline, self.subfolder)
t5xxl_filename = "t5xxl_fp16.safetensors"
t5xxl_model_path = f"{t5xxl_model_dir}/{t5xxl_filename}"
if not os.path.exists(t5xxl_model_path):
hf_hub_download(
repo_id=self.path,
filename=t5xxl_filename,
local_dir=load.get_checkpoint_dir(self.framework_model_dir, self.version, self.pipeline, ""),
subfolder=self.subfolder,
)
with safe_open(t5xxl_model_path, framework="pt", device=self.device) as f:
dtype = torch.float16 if self.fp16 else torch.float32
model = T5XXLModel(self.T5_CONFIG, device=self.device, dtype=dtype)
load_into(f, model.transformer, "", self.device, dtype)
model = optimizer.optimize_checkpoint(model, torch_inference)
return model
class CLIPVisionWithProjModel(base_model.BaseModel):
def __init__(
self,
version,
pipeline,
device,
hf_token,
verbose,
framework_model_dir,
max_batch_size=1,
subfolder="image_encoder",
):
super(CLIPVisionWithProjModel, self).__init__(
version,
pipeline,
device=device,
hf_token=hf_token,
verbose=verbose,
framework_model_dir=framework_model_dir,
max_batch_size=max_batch_size,
)
self.subfolder = subfolder
def get_model(self, torch_inference=""):
clip_model_dir = load.get_checkpoint_dir(self.framework_model_dir, self.version, self.pipeline, self.subfolder)
if not os.path.exists(clip_model_dir):
model = CLIPVisionModelWithProjection.from_pretrained(
self.path, subfolder=self.subfolder, use_safetensors=self.hf_safetensor, token=self.hf_token
).to(self.device)
model.save_pretrained(clip_model_dir)
else:
print(f"[I] Load CLIPVisionModelWithProjection model from: {clip_model_dir}")
model = CLIPVisionModelWithProjection.from_pretrained(clip_model_dir).to(self.device)
model = optimizer.optimize_checkpoint(model, torch_inference)
return model
class CLIPImageProcessorModel(base_model.BaseModel):
def __init__(
self,
version,
pipeline,
device,
hf_token,
verbose,
framework_model_dir,
max_batch_size=1,
subfolder="feature_extractor",
):
super(CLIPImageProcessorModel, self).__init__(
version,
pipeline,
device=device,
hf_token=hf_token,
verbose=verbose,
framework_model_dir=framework_model_dir,
max_batch_size=max_batch_size,
)
self.subfolder = subfolder
def get_model(self, torch_inference=""):
clip_model_dir = load.get_checkpoint_dir(self.framework_model_dir, self.version, self.pipeline, self.subfolder)
# NOTE to(device) not supported
if not os.path.exists(clip_model_dir):
model = CLIPImageProcessor.from_pretrained(
self.path, subfolder=self.subfolder, use_safetensors=self.hf_safetensor, token=self.hf_token
)
model.save_pretrained(clip_model_dir)
else:
print(f"[I] Load CLIPImageProcessor model from: {clip_model_dir}")
model = CLIPImageProcessor.from_pretrained(clip_model_dir)
model = optimizer.optimize_checkpoint(model, torch_inference)
return model
demo/Diffusion/demo_diffusion/model/diffusion_transformer.py
+326
View File
@@ -0,0 +1,326 @@
#
# SPDX-FileCopyrightText: Copyright (c) 1993-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
import os
import torch
from huggingface_hub import hf_hub_download
from safetensors import safe_open
from demo_diffusion.dynamic_import import import_from_diffusers
from demo_diffusion.model import base_model, load, optimizer
from demo_diffusion.utils_sd3.other_impls import load_into
from demo_diffusion.utils_sd3.sd3_impls import BaseModel as BaseModelSD3
# List of models to import from diffusers.models
models_to_import = [
"FluxTransformer2DModel",
]
for model in models_to_import:
globals()[model] = import_from_diffusers(model, "diffusers.models")
class SD3_MMDiTModel(base_model.BaseModel):
def __init__(
self,
version,
pipeline,
device,
hf_token,
verbose,
framework_model_dir,
shift=1.0,
fp16=False,
max_batch_size=16,
text_maxlen=77,
):
super(SD3_MMDiTModel, self).__init__(
version,
pipeline,
device=device,
hf_token=hf_token,
verbose=verbose,
framework_model_dir=framework_model_dir,
fp16=fp16,
max_batch_size=max_batch_size,
text_maxlen=text_maxlen,
)
self.subfolder = "sd3"
self.mmdit_dim = 16
self.shift = shift
self.xB = 2
def get_model(self, torch_inference=""):
sd3_model_dir = load.get_checkpoint_dir(self.framework_model_dir, self.version, self.pipeline, self.subfolder)
sd3_filename = "sd3_medium.safetensors"
sd3_model_path = f"{sd3_model_dir}/{sd3_filename}"
if not os.path.exists(sd3_model_path):
hf_hub_download(repo_id=self.path, filename=sd3_filename, local_dir=sd3_model_dir)
with safe_open(sd3_model_path, framework="pt", device=self.device) as f:
model = BaseModelSD3(
shift=self.shift, file=f, prefix="model.diffusion_model.", device=self.device, dtype=torch.float16
).eval()
load_into(f, model, "model.", self.device, torch.float16)
model = optimizer.optimize_checkpoint(model, torch_inference)
return model
def get_input_names(self):
return ["sample", "sigma", "c_crossattn", "y"]
def get_output_names(self):
return ["latent"]
def get_dynamic_axes(self):
xB = "2B" if self.xB == 2 else "B"
return {
"sample": {0: xB, 2: "H", 3: "W"},
"sigma": {0: xB},
"c_crossattn": {0: xB},
"y": {0: xB},
"latent": {0: xB, 2: "H", 3: "W"},
}
def get_input_profile(self, batch_size, image_height, image_width, static_batch, static_shape):
latent_height, latent_width = self.check_dims(batch_size, image_height, image_width)
min_batch, max_batch, _, _, _, _, min_latent_height, max_latent_height, min_latent_width, max_latent_width = (
self.get_minmax_dims(batch_size, image_height, image_width, static_batch, static_shape)
)
return {
"sample": [
(self.xB * min_batch, self.mmdit_dim, min_latent_height, min_latent_width),
(self.xB * batch_size, self.mmdit_dim, latent_height, latent_width),
(self.xB * max_batch, self.mmdit_dim, max_latent_height, max_latent_width),
],
"sigma": [(self.xB * min_batch,), (self.xB * batch_size,), (self.xB * max_batch,)],
"c_crossattn": [
(self.xB * min_batch, 154, 4096),
(self.xB * batch_size, 154, 4096),
(self.xB * max_batch, 154, 4096),
],
"y": [(self.xB * min_batch, 2048), (self.xB * batch_size, 2048), (self.xB * max_batch, 2048)],
}
def get_shape_dict(self, batch_size, image_height, image_width):
latent_height, latent_width = self.check_dims(batch_size, image_height, image_width)
return {
"sample": (self.xB * batch_size, self.mmdit_dim, latent_height, latent_width),
"sigma": (self.xB * batch_size,),
"c_crossattn": (self.xB * batch_size, 154, 4096),
"y": (self.xB * batch_size, 2048),
"latent": (self.xB * batch_size, self.mmdit_dim, latent_height, latent_width),
}
def get_sample_input(self, batch_size, image_height, image_width, static_shape):
latent_height, latent_width = self.check_dims(batch_size, image_height, image_width)
dtype = torch.float16 if self.fp16 else torch.float32
return (
torch.randn(batch_size, self.mmdit_dim, latent_height, latent_width, dtype=dtype, device=self.device),
torch.randn(batch_size, dtype=dtype, device=self.device),
{
"c_crossattn": torch.randn(batch_size, 154, 4096, dtype=dtype, device=self.device),
"y": torch.randn(batch_size, 2048, dtype=dtype, device=self.device),
},
)
class FluxTransformerModel(base_model.BaseModel):
def __init__(
self,
version,
pipeline,
device,
hf_token,
verbose,
framework_model_dir,
fp16=False,
tf32=False,
int8=False,
fp8=False,
bf16=False,
max_batch_size=16,
text_maxlen=77,
build_strongly_typed=False,
weight_streaming=False,
weight_streaming_budget_percentage=None,
):
super(FluxTransformerModel, self).__init__(
version,
pipeline,
device=device,
hf_token=hf_token,
verbose=verbose,
framework_model_dir=framework_model_dir,
fp16=fp16,
tf32=tf32,
int8=int8,
fp8=fp8,
bf16=bf16,
max_batch_size=max_batch_size,
text_maxlen=text_maxlen,
)
self.subfolder = "transformer"
self.transformer_model_dir = load.get_checkpoint_dir(
self.framework_model_dir, self.version, self.pipeline, self.subfolder
)
if not os.path.exists(self.transformer_model_dir):
self.config = FluxTransformer2DModel.load_config(self.path, subfolder=self.subfolder, token=self.hf_token)
else:
print(f"[I] Load FluxTransformer2DModel config from: {self.transformer_model_dir}")
self.config = FluxTransformer2DModel.load_config(self.transformer_model_dir)
self.build_strongly_typed = build_strongly_typed
self.weight_streaming = weight_streaming
self.weight_streaming_budget_percentage = weight_streaming_budget_percentage
self.out_channels = self.config.get("out_channels") or self.config["in_channels"]
def get_model(self, torch_inference=""):
model_opts = (
{"torch_dtype": torch.float16} if self.fp16 else {"torch_dtype": torch.bfloat16} if self.bf16 else {}
)
if not load.is_model_cached(self.transformer_model_dir, model_opts, self.hf_safetensor):
model = FluxTransformer2DModel.from_pretrained(
self.path,
subfolder=self.subfolder,
use_safetensors=self.hf_safetensor,
token=self.hf_token,
**model_opts,
).to(self.device)
model.save_pretrained(self.transformer_model_dir, **model_opts)
else:
print(f"[I] Load FluxTransformer2DModel model from: {self.transformer_model_dir}")
model = FluxTransformer2DModel.from_pretrained(self.transformer_model_dir, **model_opts).to(self.device)
if torch_inference:
model.to(memory_format=torch.channels_last)
model = optimizer.optimize_checkpoint(model, torch_inference)
return model
def get_input_names(self):
return [
"hidden_states",
"encoder_hidden_states",
"pooled_projections",
"timestep",
"img_ids",
"txt_ids",
"guidance",
]
def get_output_names(self):
return ["latent"]
def get_dynamic_axes(self):
dynamic_axes = {
"hidden_states": {0: "B", 1: "latent_dim"},
"encoder_hidden_states": {0: "B"},
"pooled_projections": {0: "B"},
"timestep": {0: "B"},
"img_ids": {0: "latent_dim"},
}
if self.config["guidance_embeds"]:
dynamic_axes["guidance"] = {0: "B"}
return dynamic_axes
def get_input_profile(self, batch_size, image_height, image_width, static_batch, static_shape):
latent_height, latent_width = self.check_dims(batch_size, image_height, image_width)
(
min_batch,
max_batch,
min_image_height,
max_image_height,
min_image_width,
max_image_width,
min_latent_height,
max_latent_height,
min_latent_width,
max_latent_width,
) = self.get_minmax_dims(batch_size, image_height, image_width, static_batch, static_shape)
input_profile = {
"hidden_states": [
(min_batch, (min_latent_height // 2) * (min_latent_width // 2), self.config["in_channels"]),
(batch_size, (latent_height // 2) * (latent_width // 2), self.config["in_channels"]),
(max_batch, (max_latent_height // 2) * (max_latent_width // 2), self.config["in_channels"]),
],
"encoder_hidden_states": [
(min_batch, self.text_maxlen, self.config["joint_attention_dim"]),
(batch_size, self.text_maxlen, self.config["joint_attention_dim"]),
(max_batch, self.text_maxlen, self.config["joint_attention_dim"]),
],
"pooled_projections": [
(min_batch, self.config["pooled_projection_dim"]),
(batch_size, self.config["pooled_projection_dim"]),
(max_batch, self.config["pooled_projection_dim"]),
],
"timestep": [(min_batch,), (batch_size,), (max_batch,)],
"img_ids": [
((min_latent_height // 2) * (min_latent_width // 2), 3),
((latent_height // 2) * (latent_width // 2), 3),
((max_latent_height // 2) * (max_latent_width // 2), 3),
],
"txt_ids": [(self.text_maxlen, 3), (self.text_maxlen, 3), (self.text_maxlen, 3)],
}
if self.config["guidance_embeds"]:
input_profile["guidance"] = [(min_batch,), (batch_size,), (max_batch,)]
return input_profile
def get_shape_dict(self, batch_size, image_height, image_width):
latent_height, latent_width = self.check_dims(batch_size, image_height, image_width)
shape_dict = {
"hidden_states": (batch_size, (latent_height // 2) * (latent_width // 2), self.config["in_channels"]),
"encoder_hidden_states": (batch_size, self.text_maxlen, self.config["joint_attention_dim"]),
"pooled_projections": (batch_size, self.config["pooled_projection_dim"]),
"timestep": (batch_size,),
"img_ids": ((latent_height // 2) * (latent_width // 2), 3),
"txt_ids": (self.text_maxlen, 3),
"latent": (batch_size, (latent_height // 2) * (latent_width // 2), self.out_channels),
}
if self.config["guidance_embeds"]:
shape_dict["guidance"] = (batch_size,)
return shape_dict
def get_sample_input(self, batch_size, image_height, image_width, static_shape):
latent_height, latent_width = self.check_dims(batch_size, image_height, image_width)
dtype = torch.float32
assert not (self.fp16 and self.bf16), "fp16 and bf16 cannot be enabled simultaneously"
tensor_dtype = torch.bfloat16 if self.bf16 else (torch.float16 if self.fp16 else torch.float32)
sample_input = (
torch.randn(
batch_size,
(latent_height // 2) * (latent_width // 2),
self.config["in_channels"],
dtype=tensor_dtype,
device=self.device,
),
torch.randn(
batch_size, self.text_maxlen, self.config["joint_attention_dim"], dtype=tensor_dtype, device=self.device
),
torch.randn(batch_size, self.config["pooled_projection_dim"], dtype=tensor_dtype, device=self.device),
torch.tensor([1.0] * batch_size, dtype=tensor_dtype, device=self.device),
torch.randn((latent_height // 2) * (latent_width // 2), 3, dtype=dtype, device=self.device),
torch.randn(self.text_maxlen, 3, dtype=dtype, device=self.device),
{},
)
if self.config["guidance_embeds"]:
sample_input[-1]["guidance"] = torch.tensor([1.0] * batch_size, dtype=dtype, device=self.device)
return sample_input
def optimize(self, onnx_graph):
if self.fp8:
return super().optimize(onnx_graph)
if self.int8:
return super().optimize(onnx_graph, fuse_mha_qkv_int8=True)
return super().optimize(onnx_graph)
demo/Diffusion/demo_diffusion/model/gan.py
+143
View File
@@ -0,0 +1,143 @@
#
# SPDX-FileCopyrightText: Copyright (c) 1993-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
import torch
from diffusers.pipelines.wuerstchen import PaellaVQModel
from demo_diffusion.model import base_model, load, optimizer
class VQGANModel(base_model.BaseModel):
def __init__(
self,
version,
pipeline,
device,
hf_token,
verbose,
framework_model_dir,
fp16=False,
bf16=False,
max_batch_size=16,
compression_factor=42,
latent_dim_scale=10.67,
scale_factor=0.3764,
):
super(VQGANModel, self).__init__(
version,
pipeline,
device=device,
hf_token=hf_token,
verbose=verbose,
framework_model_dir=framework_model_dir,
fp16=fp16,
bf16=bf16,
max_batch_size=max_batch_size,
compression_factor=compression_factor,
)
self.subfolder = "vqgan"
self.latent_dim_scale = latent_dim_scale
self.scale_factor = scale_factor
def get_model(self, torch_inference=""):
model_opts = {"variant": "bf16", "torch_dtype": torch.bfloat16} if self.bf16 else {}
vqgan_model_dir = load.get_checkpoint_dir(self.framework_model_dir, self.version, self.pipeline, self.subfolder)
if not load.is_model_cached(vqgan_model_dir, model_opts, self.hf_safetensor, model_name="model"):
model = PaellaVQModel.from_pretrained(
self.path,
subfolder=self.subfolder,
use_safetensors=self.hf_safetensor,
token=self.hf_token,
**model_opts,
).to(self.device)
model.save_pretrained(vqgan_model_dir, **model_opts)
else:
print(f"[I] Load VQGAN pytorch model from: {vqgan_model_dir}")
model = PaellaVQModel.from_pretrained(vqgan_model_dir, **model_opts).to(self.device)
model.forward = model.decode
model = optimizer.optimize_checkpoint(model, torch_inference)
return model
def get_input_names(self):
return ["latent"]
def get_output_names(self):
return ["images"]
def get_dynamic_axes(self):
return {"latent": {0: "B", 2: "H", 3: "W"}, "images": {0: "B", 2: "8H", 3: "8W"}}
def get_input_profile(self, batch_size, image_height, image_width, static_batch, static_shape):
latent_height, latent_width = self.check_dims(batch_size, image_height, image_width)
min_batch, max_batch, _, _, _, _, min_latent_height, max_latent_height, min_latent_width, max_latent_width = (
self.get_minmax_dims(batch_size, image_height, image_width, static_batch, static_shape)
)
return {
"latent": [
(min_batch, 4, min_latent_height, min_latent_width),
(batch_size, 4, latent_height, latent_width),
(max_batch, 4, max_latent_height, max_latent_width),
]
}
def get_shape_dict(self, batch_size, image_height, image_width):
latent_height, latent_width = self.check_dims(batch_size, image_height, image_width)
return {
"latent": (batch_size, 4, latent_height, latent_width),
"images": (batch_size, 3, image_height, image_width),
}
def get_sample_input(self, batch_size, image_height, image_width, static_shape):
latent_height, latent_width = self.check_dims(batch_size, image_height, image_width)
dtype = torch.float16 if self.fp16 else torch.bfloat16 if self.bf16 else torch.float32
return torch.randn(batch_size, 4, latent_height, latent_width, dtype=dtype, device=self.device)
def check_dims(self, batch_size, image_height, image_width):
latent_height, latent_width = super().check_dims(batch_size, image_height, image_width)
latent_height = int(latent_height * self.latent_dim_scale)
latent_width = int(latent_width * self.latent_dim_scale)
return (latent_height, latent_width)
def get_minmax_dims(self, batch_size, image_height, image_width, static_batch, static_shape):
(
min_batch,
max_batch,
min_image_height,
max_image_height,
min_image_width,
max_image_width,
min_latent_height,
max_latent_height,
min_latent_width,
max_latent_width,
) = super().get_minmax_dims(batch_size, image_height, image_width, static_batch, static_shape)
min_latent_height = int(min_latent_height * self.latent_dim_scale)
min_latent_width = int(min_latent_width * self.latent_dim_scale)
max_latent_height = int(max_latent_height * self.latent_dim_scale)
max_latent_width = int(max_latent_width * self.latent_dim_scale)
return (
min_batch,
max_batch,
min_image_height,
max_image_height,
min_image_width,
max_image_width,
min_latent_height,
max_latent_height,
min_latent_width,
max_latent_width,
)
demo/Diffusion/demo_diffusion/model/load.py
+115
View File
@@ -0,0 +1,115 @@
#
# SPDX-FileCopyrightText: Copyright (c) 1993-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
"""
Functions for loading models.
"""
from __future__ import annotations
import gc
import glob
import os
import sys
from typing import List, Optional
import torch
import onnx
def onnx_graph_needs_external_data(onnx_graph: onnx.ModelProto) -> bool:
"""Return true if ONNX graph needs to store external data."""
if sys.platform == "win32":
# ByteSize is broken (wraps around) on Windows, so always assume external data is needed.
return True
else:
TWO_GIGABYTES = 2147483648
return onnx_graph.ByteSize() > TWO_GIGABYTES
def get_path(version: str, pipeline: "pipeline.DiffusionPipeline", controlnets: Optional[List[str]] = None) -> str:
"""Return the relative path to the model files directory."""
if controlnets is not None:
if version == "xl-1.0":
return ["diffusers/controlnet-canny-sdxl-1.0"]
return ["lllyasviel/sd-controlnet-" + modality for modality in controlnets]
if version in ("1.4", "1.5") and pipeline.is_inpaint():
return "benjamin-paine/stable-diffusion-v1-5-inpainting"
elif version == "1.4":
return "CompVis/stable-diffusion-v1-4"
elif version == "1.5":
return "KiwiXR/stable-diffusion-v1-5"
elif version == "dreamshaper-7":
return "Lykon/dreamshaper-7"
elif version in ("2.0-base", "2.0") and pipeline.is_inpaint():
return "stabilityai/stable-diffusion-2-inpainting"
elif version == "2.0-base":
return "stabilityai/stable-diffusion-2-base"
elif version == "2.0":
return "stabilityai/stable-diffusion-2"
elif version == "2.1-base":
return "stabilityai/stable-diffusion-2-1-base"
elif version == "2.1":
return "stabilityai/stable-diffusion-2-1"
elif version == "xl-1.0" and pipeline.is_sd_xl_base():
return "stabilityai/stable-diffusion-xl-base-1.0"
elif version == "xl-1.0" and pipeline.is_sd_xl_refiner():
return "stabilityai/stable-diffusion-xl-refiner-1.0"
# TODO SDXL turbo with refiner
elif version == "xl-turbo" and pipeline.is_sd_xl_base():
return "stabilityai/sdxl-turbo"
elif version == "sd3":
return "stabilityai/stable-diffusion-3-medium"
elif version == "svd-xt-1.1" and pipeline.is_img2vid():
return "stabilityai/stable-video-diffusion-img2vid-xt-1-1"
elif version == "cascade":
if pipeline.is_cascade_decoder():
return "stabilityai/stable-cascade"
else:
return "stabilityai/stable-cascade-prior"
elif version == "flux.1-dev":
return "black-forest-labs/FLUX.1-dev"
elif version == "flux.1-schnell":
return "black-forest-labs/FLUX.1-schnell"
elif version == "flux.1-dev-canny":
return "black-forest-labs/FLUX.1-Canny-dev"
elif version == "flux.1-dev-depth":
return "black-forest-labs/FLUX.1-Depth-dev"
else:
raise ValueError(f"Unsupported version {version} + pipeline {pipeline.name}")
# FIXME serialization not supported for torch.compile
def get_checkpoint_dir(framework_model_dir: str, version: str, pipeline: str, subfolder: str) -> str:
"""Return the path to the torch model checkpoint directory."""
return os.path.join(framework_model_dir, version, pipeline, subfolder)
def is_model_cached(model_dir, model_opts, hf_safetensor, model_name="diffusion_pytorch_model") -> bool:
"""Return True if model was cached."""
variant = "." + model_opts.get("variant") if "variant" in model_opts else ""
suffix = ".safetensors" if hf_safetensor else ".bin"
# WAR with * for larger models that are split into multiple smaller ckpt files
model_file = model_name + variant + "*" + suffix
return bool(glob.glob(os.path.join(model_dir, model_file)))
def unload_torch_model(model):
if model:
del model
torch.cuda.empty_cache()
gc.collect()
demo/Diffusion/demo_diffusion/model/lora.py
+39
View File
@@ -0,0 +1,39 @@
#
# SPDX-FileCopyrightText: Copyright (c) 1993-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
from diffusers.loaders import StableDiffusionLoraLoaderMixin
class LoraLoader(StableDiffusionLoraLoaderMixin):
def __init__(self, paths, weights, scale):
self.paths = paths
self.weights = weights
self.scale = scale
def merge_loras(model, lora_loader):
paths, weights, scale = lora_loader.paths, lora_loader.weights, lora_loader.scale
for i, path in enumerate(paths):
print(f"[I] Loading LoRA: {path}, weight {weights[i]}")
state_dict, network_alphas = lora_loader.lora_state_dict(path, unet_config=model.config)
lora_loader.load_lora_into_unet(state_dict, network_alphas=network_alphas, unet=model, adapter_name=path)
model.set_adapters(paths, weights=weights)
# NOTE: fuse_lora an experimental API in Diffusers
model.fuse_lora(adapter_names=paths, lora_scale=scale)
model.unload_lora()
return model
demo/Diffusion/demo_diffusion/model/optimizer.py
+204
View File
@@ -0,0 +1,204 @@
#
# SPDX-FileCopyrightText: Copyright (c) 1993-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
import os
import re
import tempfile
import onnx_graphsurgeon as gs
import torch
from onnxmltools.utils.float16_converter import convert_float_to_float16
from polygraphy.backend.onnx.loader import fold_constants
import onnx
from demo_diffusion.model import load
from demo_diffusion.utils_modelopt import (
cast_fp8_mha_io,
cast_resize_io,
convert_fp16_io,
convert_zp_fp8,
)
from onnx import shape_inference
# FIXME update callsites after serialization support for torch.compile is added
TORCH_INFERENCE_MODELS = ["default", "reduce-overhead", "max-autotune"]
def optimize_checkpoint(model, torch_inference: str):
"""Optimize a torch model checkpoint using torch.compile."""
if not torch_inference or torch_inference == "eager":
return model
assert torch_inference in TORCH_INFERENCE_MODELS
return torch.compile(model, mode=torch_inference, dynamic=False, fullgraph=False)
class Optimizer:
def __init__(self, onnx_graph, verbose=False):
self.graph = gs.import_onnx(onnx_graph)
self.verbose = verbose
def info(self, prefix):
if self.verbose:
print(
f"{prefix} .. {len(self.graph.nodes)} nodes, {len(self.graph.tensors().keys())} tensors, {len(self.graph.inputs)} inputs, {len(self.graph.outputs)} outputs"
)
def cleanup(self, return_onnx=False):
self.graph.cleanup().toposort()
return gs.export_onnx(self.graph) if return_onnx else self.graph
def select_outputs(self, keep, names=None):
self.graph.outputs = [self.graph.outputs[o] for o in keep]
if names:
for i, name in enumerate(names):
self.graph.outputs[i].name = name
def fold_constants(self, return_onnx=False):
onnx_graph = fold_constants(gs.export_onnx(self.graph), allow_onnxruntime_shape_inference=True)
self.graph = gs.import_onnx(onnx_graph)
if return_onnx:
return onnx_graph
def infer_shapes(self, return_onnx=False):
onnx_graph = gs.export_onnx(self.graph)
if load.onnx_graph_needs_external_data(onnx_graph):
temp_dir = tempfile.TemporaryDirectory().name
os.makedirs(temp_dir, exist_ok=True)
onnx_orig_path = os.path.join(temp_dir, "model.onnx")
onnx_inferred_path = os.path.join(temp_dir, "inferred.onnx")
onnx.save_model(
onnx_graph,
onnx_orig_path,
save_as_external_data=True,
all_tensors_to_one_file=True,
convert_attribute=False,
)
onnx.shape_inference.infer_shapes_path(onnx_orig_path, onnx_inferred_path)
onnx_graph = onnx.load(onnx_inferred_path)
else:
onnx_graph = shape_inference.infer_shapes(onnx_graph)
self.graph = gs.import_onnx(onnx_graph)
if return_onnx:
return onnx_graph
def clip_add_hidden_states(self, hidden_layer_offset, return_onnx=False):
hidden_layers = -1
onnx_graph = gs.export_onnx(self.graph)
for i in range(len(onnx_graph.graph.node)):
for j in range(len(onnx_graph.graph.node[i].output)):
name = onnx_graph.graph.node[i].output[j]
if "layers" in name:
hidden_layers = max(int(name.split(".")[1].split("/")[0]), hidden_layers)
for i in range(len(onnx_graph.graph.node)):
for j in range(len(onnx_graph.graph.node[i].output)):
if onnx_graph.graph.node[i].output[j] == "/text_model/encoder/layers.{}/Add_1_output_0".format(
hidden_layers + hidden_layer_offset
):
onnx_graph.graph.node[i].output[j] = "hidden_states"
for j in range(len(onnx_graph.graph.node[i].input)):
if onnx_graph.graph.node[i].input[j] == "/text_model/encoder/layers.{}/Add_1_output_0".format(
hidden_layers + hidden_layer_offset
):
onnx_graph.graph.node[i].input[j] = "hidden_states"
if return_onnx:
return onnx_graph
def fuse_mha_qkv_int8_sq(self):
tensors = self.graph.tensors()
keys = tensors.keys()
# mha : fuse QKV QDQ nodes
# mhca : fuse KV QDQ nodes
q_pat = (
"/down_blocks.\\d+/attentions.\\d+/transformer_blocks"
".\\d+/attn\\d+/to_q/input_quantizer/DequantizeLinear_output_0"
)
k_pat = (
"/down_blocks.\\d+/attentions.\\d+/transformer_blocks"
".\\d+/attn\\d+/to_k/input_quantizer/DequantizeLinear_output_0"
)
v_pat = (
"/down_blocks.\\d+/attentions.\\d+/transformer_blocks"
".\\d+/attn\\d+/to_v/input_quantizer/DequantizeLinear_output_0"
)
qs = list(
sorted(
map(
lambda x: x.group(0), # type: ignore
filter(lambda x: x is not None, [re.match(q_pat, key) for key in keys]),
)
)
)
ks = list(
sorted(
map(
lambda x: x.group(0), # type: ignore
filter(lambda x: x is not None, [re.match(k_pat, key) for key in keys]),
)
)
)
vs = list(
sorted(
map(
lambda x: x.group(0), # type: ignore
filter(lambda x: x is not None, [re.match(v_pat, key) for key in keys]),
)
)
)
removed = 0
assert len(qs) == len(ks) == len(vs), "Failed to collect tensors"
for q, k, v in zip(qs, ks, vs):
is_mha = all(["attn1" in tensor for tensor in [q, k, v]])
is_mhca = all(["attn2" in tensor for tensor in [q, k, v]])
assert (is_mha or is_mhca) and (not (is_mha and is_mhca))
if is_mha:
tensors[k].outputs[0].inputs[0] = tensors[q]
tensors[v].outputs[0].inputs[0] = tensors[q]
del tensors[k]
del tensors[v]
removed += 2
else: # is_mhca
tensors[k].outputs[0].inputs[0] = tensors[v]
del tensors[k]
removed += 1
print(f"Removed {removed} QDQ nodes")
return removed # expected 72 for L2.5
def modify_fp8_graph(self, is_fp16_io=True):
onnx_graph = gs.export_onnx(self.graph)
# Convert INT8 Zero to FP8.
onnx_graph = convert_zp_fp8(onnx_graph)
# Convert weights and activations to FP16 and insert Cast nodes in FP8 MHA.
onnx_graph = convert_float_to_float16(onnx_graph, keep_io_types=True, disable_shape_infer=True)
self.graph = gs.import_onnx(onnx_graph)
# Add cast nodes to Resize I/O.
cast_resize_io(self.graph)
# Convert model inputs and outputs to fp16 I/O.
if is_fp16_io:
convert_fp16_io(self.graph)
# Add cast nodes to MHA's BMM1 and BMM2's I/O.
cast_fp8_mha_io(self.graph)
def flux_convert_rope_weight_type(self):
for node in self.graph.nodes:
if node.op == "Einsum":
print(f"Fixed RoPE (Rotary Position Embedding) weight type: {node.name}")
return gs.export_onnx(self.graph)
demo/Diffusion/demo_diffusion/model/scheduler.py
+33
View File
@@ -0,0 +1,33 @@
#
# SPDX-FileCopyrightText: Copyright (c) 1993-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
import os
from demo_diffusion.model import load
def make_scheduler(cls, version, pipeline, hf_token, framework_model_dir, subfolder="scheduler"):
scheduler_dir = os.path.join(
framework_model_dir, version, pipeline.name, next(iter({cls.__name__})).lower(), subfolder
)
if not os.path.exists(scheduler_dir):
scheduler = cls.from_pretrained(load.get_path(version, pipeline), subfolder=subfolder, token=hf_token)
scheduler.save_pretrained(scheduler_dir)
else:
print(f"[I] Load Scheduler {cls.__name__} from: {scheduler_dir}")
scheduler = cls.from_pretrained(scheduler_dir)
return scheduler
demo/Diffusion/demo_diffusion/model/t5.py
+121
View File
@@ -0,0 +1,121 @@
#
# SPDX-FileCopyrightText: Copyright (c) 1993-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
import os
import torch
from transformers import (
AutoConfig,
T5EncoderModel,
)
from demo_diffusion.model import base_model, load, optimizer
class T5Model(base_model.BaseModel):
def __init__(
self,
version,
pipeline,
device,
hf_token,
verbose,
framework_model_dir,
max_batch_size,
fp16=False,
tf32=False,
bf16=False,
subfolder="text_encoder",
text_maxlen=512,
build_strongly_typed=False,
weight_streaming=False,
weight_streaming_budget_percentage=None,
):
super(T5Model, self).__init__(
version,
pipeline,
device=device,
hf_token=hf_token,
verbose=verbose,
framework_model_dir=framework_model_dir,
fp16=fp16,
tf32=tf32,
bf16=bf16,
max_batch_size=max_batch_size,
text_maxlen=text_maxlen,
)
self.subfolder = subfolder
self.t5_model_dir = load.get_checkpoint_dir(
self.framework_model_dir, self.version, self.pipeline, self.subfolder
)
if not os.path.exists(self.t5_model_dir):
self.config = AutoConfig.from_pretrained(self.path, subfolder=self.subfolder, token=self.hf_token)
else:
print(f"[I] Load T5Encoder Config from: {self.t5_model_dir}")
self.config = AutoConfig.from_pretrained(self.t5_model_dir)
self.build_strongly_typed = build_strongly_typed
self.weight_streaming = weight_streaming
self.weight_streaming_budget_percentage = weight_streaming_budget_percentage
def get_model(self, torch_inference=""):
model_opts = (
{"torch_dtype": torch.float16} if self.fp16 else {"torch_dtype": torch.bfloat16} if self.bf16 else {}
)
if not load.is_model_cached(self.t5_model_dir, model_opts, self.hf_safetensor, model_name="model"):
model = T5EncoderModel.from_pretrained(
self.path,
subfolder=self.subfolder,
use_safetensors=self.hf_safetensor,
token=self.hf_token,
**model_opts,
).to(self.device)
model.save_pretrained(self.t5_model_dir, **model_opts)
else:
print(f"[I] Load T5EncoderModel model from: {self.t5_model_dir}")
model = T5EncoderModel.from_pretrained(self.t5_model_dir, **model_opts).to(self.device)
model = optimizer.optimize_checkpoint(model, torch_inference)
return model
def get_input_names(self):
return ["input_ids"]
def get_output_names(self):
return ["text_embeddings"]
def get_dynamic_axes(self):
return {"input_ids": {0: "B"}, "text_embeddings": {0: "B"}}
def get_input_profile(self, batch_size, image_height, image_width, static_batch, static_shape):
self.check_dims(batch_size, image_height, image_width)
min_batch, max_batch, _, _, _, _, _, _, _, _ = self.get_minmax_dims(
batch_size, image_height, image_width, static_batch, static_shape
)
return {
"input_ids": [(min_batch, self.text_maxlen), (batch_size, self.text_maxlen), (max_batch, self.text_maxlen)]
}
def get_shape_dict(self, batch_size, image_height, image_width):
self.check_dims(batch_size, image_height, image_width)
output = {
"input_ids": (batch_size, self.text_maxlen),
"text_embeddings": (batch_size, self.text_maxlen, self.config.d_model),
}
return output
def get_sample_input(self, batch_size, image_height, image_width, static_shape):
self.check_dims(batch_size, image_height, image_width)
return torch.zeros(batch_size, self.text_maxlen, dtype=torch.int32, device=self.device)
demo/Diffusion/demo_diffusion/model/tokenizer.py
+46
View File
@@ -0,0 +1,46 @@
#
# SPDX-FileCopyrightText: Copyright (c) 1993-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
import os
from transformers import (
CLIPTokenizer,
T5TokenizerFast,
)
from demo_diffusion.model import load
def make_tokenizer(version, pipeline, hf_token, framework_model_dir, subfolder="tokenizer", tokenizer_type="clip"):
if tokenizer_type == "clip":
tokenizer_class = CLIPTokenizer
elif tokenizer_type == "t5":
tokenizer_class = T5TokenizerFast
else:
raise ValueError(
f"Unsupported tokenizer_type {tokenizer_type}. Only tokenizer_type clip and t5 are currently supported"
)
tokenizer_model_dir = load.get_checkpoint_dir(framework_model_dir, version, pipeline.name, subfolder)
if not os.path.exists(tokenizer_model_dir):
model = tokenizer_class.from_pretrained(
load.get_path(version, pipeline), subfolder=subfolder, use_safetensors=pipeline.is_sd_xl(), token=hf_token
)
model.save_pretrained(tokenizer_model_dir)
else:
print(f"[I] Load {tokenizer_class.__name__} model from: {tokenizer_model_dir}")
model = tokenizer_class.from_pretrained(tokenizer_model_dir)
return model
demo/Diffusion/demo_diffusion/model/unet.py
+883
View File
@@ -0,0 +1,883 @@
#
# SPDX-FileCopyrightText: Copyright (c) 1993-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
"""
Model definitions for UNet models.
"""
import torch
from demo_diffusion.dynamic_import import import_from_diffusers
from demo_diffusion.model import base_model, load, optimizer
# List of models to import from diffusers.models
models_to_import = [
"ControlNetModel",
"UNet2DConditionModel",
"UNetSpatioTemporalConditionModel",
"StableCascadeUNet",
]
for model in models_to_import:
globals()[model] = import_from_diffusers(model, "diffusers.models")
def get_unet_embedding_dim(version, pipeline):
if version in ("1.4", "1.5", "dreamshaper-7"):
return 768
elif version in ("2.0", "2.0-base", "2.1", "2.1-base"):
return 1024
elif version in ("xl-1.0", "xl-turbo") and pipeline.is_sd_xl_base():
return 2048
elif version in ("cascade"):
return 1280
elif version in ("xl-1.0", "xl-turbo") and pipeline.is_sd_xl_refiner():
return 1280
elif pipeline.is_img2vid():
return 1024
else:
raise ValueError(f"Invalid version {version} + pipeline {pipeline}")
class UNet2DConditionControlNetModel(torch.nn.Module):
def __init__(self, unet, controlnets) -> None:
super().__init__()
self.unet = unet
self.controlnets = controlnets
def forward(self, sample, timestep, encoder_hidden_states, images, controlnet_scales, added_cond_kwargs=None):
for i, (image, conditioning_scale, controlnet) in enumerate(zip(images, controlnet_scales, self.controlnets)):
down_samples, mid_sample = controlnet(
sample,
timestep,
encoder_hidden_states=encoder_hidden_states,
controlnet_cond=image,
return_dict=False,
added_cond_kwargs=added_cond_kwargs,
)
down_samples = [down_sample * conditioning_scale for down_sample in down_samples]
mid_sample *= conditioning_scale
# merge samples
if i == 0:
down_block_res_samples, mid_block_res_sample = down_samples, mid_sample
else:
down_block_res_samples = [
samples_prev + samples_curr
for samples_prev, samples_curr in zip(down_block_res_samples, down_samples)
]
mid_block_res_sample += mid_sample
noise_pred = self.unet(
sample,
timestep,
encoder_hidden_states=encoder_hidden_states,
down_block_additional_residuals=down_block_res_samples,
mid_block_additional_residual=mid_block_res_sample,
added_cond_kwargs=added_cond_kwargs,
)
return noise_pred
class UNetModel(base_model.BaseModel):
def __init__(
self,
version,
pipeline,
device,
hf_token,
verbose,
framework_model_dir,
fp16=False,
int8=False,
fp8=False,
max_batch_size=16,
text_maxlen=77,
controlnets=None,
do_classifier_free_guidance=False,
):
super(UNetModel, self).__init__(
version,
pipeline,
device=device,
hf_token=hf_token,
verbose=verbose,
framework_model_dir=framework_model_dir,
fp16=fp16,
int8=int8,
fp8=fp8,
max_batch_size=max_batch_size,
text_maxlen=text_maxlen,
embedding_dim=get_unet_embedding_dim(version, pipeline),
)
self.subfolder = "unet"
self.controlnets = load.get_path(version, pipeline, controlnets) if controlnets else None
self.unet_dim = 9 if pipeline.is_inpaint() else 4
self.xB = 2 if do_classifier_free_guidance else 1 # batch multiplier
def get_model(self, torch_inference=""):
model_opts = {"variant": "fp16", "torch_dtype": torch.float16} if self.fp16 else {}
if self.controlnets:
unet_model = UNet2DConditionModel.from_pretrained(
self.path,
subfolder=self.subfolder,
use_safetensors=self.hf_safetensor,
token=self.hf_token,
**model_opts,
).to(self.device)
cnet_model_opts = {"torch_dtype": torch.float16} if self.fp16 else {}
controlnets = torch.nn.ModuleList(
[ControlNetModel.from_pretrained(path, **cnet_model_opts).to(self.device) for path in self.controlnets]
)
# FIXME - cache UNet2DConditionControlNetModel
model = UNet2DConditionControlNetModel(unet_model, controlnets)
else:
unet_model_dir = load.get_checkpoint_dir(
self.framework_model_dir, self.version, self.pipeline, self.subfolder
)
if not load.is_model_cached(unet_model_dir, model_opts, self.hf_safetensor):
model = UNet2DConditionModel.from_pretrained(
self.path,
subfolder=self.subfolder,
use_safetensors=self.hf_safetensor,
token=self.hf_token,
**model_opts,
).to(self.device)
model.save_pretrained(unet_model_dir, **model_opts)
else:
print(f"[I] Load UNet2DConditionModel model from: {unet_model_dir}")
model = UNet2DConditionModel.from_pretrained(unet_model_dir, **model_opts).to(self.device)
if torch_inference:
model.to(memory_format=torch.channels_last)
model = optimizer.optimize_checkpoint(model, torch_inference)
return model
def get_input_names(self):
if self.controlnets is None:
return ["sample", "timestep", "encoder_hidden_states"]
else:
return ["sample", "timestep", "encoder_hidden_states", "images", "controlnet_scales"]
def get_output_names(self):
return ["latent"]
def get_dynamic_axes(self):
xB = "2B" if self.xB == 2 else "B"
if self.controlnets is None:
return {
"sample": {0: xB, 2: "H", 3: "W"},
"encoder_hidden_states": {0: xB},
"latent": {0: xB, 2: "H", 3: "W"},
}
else:
return {
"sample": {0: xB, 2: "H", 3: "W"},
"encoder_hidden_states": {0: xB},
"images": {1: xB, 3: "8H", 4: "8W"},
"latent": {0: xB, 2: "H", 3: "W"},
}
def get_input_profile(self, batch_size, image_height, image_width, static_batch, static_shape):
# WAR to enable inference for H/W that are not multiples of 16
# If building with Dynamic Shapes: ensure image height and width are not multiples of 16 for ONNX export and TensorRT engine build
if not static_shape:
image_height = image_height - 8 if image_height % 16 == 0 else image_height
image_width = image_width - 8 if image_width % 16 == 0 else image_width
latent_height, latent_width = self.check_dims(batch_size, image_height, image_width)
(
min_batch,
max_batch,
min_image_height,
max_image_height,
min_image_width,
max_image_width,
min_latent_height,
max_latent_height,
min_latent_width,
max_latent_width,
) = self.get_minmax_dims(batch_size, image_height, image_width, static_batch, static_shape)
if self.controlnets is None:
return {
"sample": [
(self.xB * min_batch, self.unet_dim, min_latent_height, min_latent_width),
(self.xB * batch_size, self.unet_dim, latent_height, latent_width),
(self.xB * max_batch, self.unet_dim, max_latent_height, max_latent_width),
],
"encoder_hidden_states": [
(self.xB * min_batch, self.text_maxlen, self.embedding_dim),
(self.xB * batch_size, self.text_maxlen, self.embedding_dim),
(self.xB * max_batch, self.text_maxlen, self.embedding_dim),
],
}
else:
return {
"sample": [
(self.xB * min_batch, self.unet_dim, min_latent_height, min_latent_width),
(self.xB * batch_size, self.unet_dim, latent_height, latent_width),
(self.xB * max_batch, self.unet_dim, max_latent_height, max_latent_width),
],
"encoder_hidden_states": [
(self.xB * min_batch, self.text_maxlen, self.embedding_dim),
(self.xB * batch_size, self.text_maxlen, self.embedding_dim),
(self.xB * max_batch, self.text_maxlen, self.embedding_dim),
],
"images": [
(len(self.controlnets), self.xB * min_batch, 3, min_image_height, min_image_width),
(len(self.controlnets), self.xB * batch_size, 3, image_height, image_width),
(len(self.controlnets), self.xB * max_batch, 3, max_image_height, max_image_width),
],
}
def get_shape_dict(self, batch_size, image_height, image_width):
latent_height, latent_width = self.check_dims(batch_size, image_height, image_width)
if self.controlnets is None:
return {
"sample": (self.xB * batch_size, self.unet_dim, latent_height, latent_width),
"encoder_hidden_states": (self.xB * batch_size, self.text_maxlen, self.embedding_dim),
"latent": (self.xB * batch_size, 4, latent_height, latent_width),
}
else:
return {
"sample": (self.xB * batch_size, self.unet_dim, latent_height, latent_width),
"encoder_hidden_states": (self.xB * batch_size, self.text_maxlen, self.embedding_dim),
"images": (len(self.controlnets), self.xB * batch_size, 3, image_height, image_width),
"latent": (self.xB * batch_size, 4, latent_height, latent_width),
}
def get_sample_input(self, batch_size, image_height, image_width, static_shape):
# WAR to enable inference for H/W that are not multiples of 16
# If building with Dynamic Shapes: ensure image height and width are not multiples of 16 for ONNX export and TensorRT engine build
if not static_shape:
image_height = image_height - 8 if image_height % 16 == 0 else image_height
image_width = image_width - 8 if image_width % 16 == 0 else image_width
latent_height, latent_width = self.check_dims(batch_size, image_height, image_width)
dtype = torch.float16 if self.fp16 else torch.float32
if self.controlnets is None:
return (
torch.randn(batch_size, self.unet_dim, latent_height, latent_width, dtype=dtype, device=self.device),
torch.tensor([1.0], dtype=dtype, device=self.device),
torch.randn(batch_size, self.text_maxlen, self.embedding_dim, dtype=dtype, device=self.device),
)
else:
return (
torch.randn(batch_size, self.unet_dim, latent_height, latent_width, dtype=dtype, device=self.device),
torch.tensor(999, dtype=dtype, device=self.device),
torch.randn(batch_size, self.text_maxlen, self.embedding_dim, dtype=dtype, device=self.device),
torch.randn(
len(self.controlnets), batch_size, 3, image_height, image_width, dtype=dtype, device=self.device
),
torch.randn(len(self.controlnets), dtype=dtype, device=self.device),
)
def optimize(self, onnx_graph):
if self.fp8:
return super().optimize(onnx_graph, modify_fp8_graph=True)
if self.int8:
return super().optimize(onnx_graph, fuse_mha_qkv_int8=True)
return super().optimize(onnx_graph)
class UNetXLModel(base_model.BaseModel):
def __init__(
self,
version,
pipeline,
device,
hf_token,
verbose,
framework_model_dir,
fp16=False,
int8=False,
fp8=False,
max_batch_size=16,
text_maxlen=77,
do_classifier_free_guidance=False,
):
super(UNetXLModel, self).__init__(
version,
pipeline,
device=device,
hf_token=hf_token,
verbose=verbose,
framework_model_dir=framework_model_dir,
fp16=fp16,
int8=int8,
fp8=fp8,
max_batch_size=max_batch_size,
text_maxlen=text_maxlen,
embedding_dim=get_unet_embedding_dim(version, pipeline),
)
self.subfolder = "unet"
self.unet_dim = 9 if pipeline.is_inpaint() else 4
self.time_dim = 5 if pipeline.is_sd_xl_refiner() else 6
self.xB = 2 if do_classifier_free_guidance else 1 # batch multiplier
def get_model(self, torch_inference=""):
model_opts = {"variant": "fp16", "torch_dtype": torch.float16} if self.fp16 else {}
unet_model_dir = load.get_checkpoint_dir(self.framework_model_dir, self.version, self.pipeline, self.subfolder)
if not load.is_model_cached(unet_model_dir, model_opts, self.hf_safetensor):
model = UNet2DConditionModel.from_pretrained(
self.path,
subfolder=self.subfolder,
use_safetensors=self.hf_safetensor,
token=self.hf_token,
**model_opts,
).to(self.device)
# Use default attention processor for ONNX export
if not torch_inference:
model.set_default_attn_processor()
model.save_pretrained(unet_model_dir, **model_opts)
else:
print(f"[I] Load UNet2DConditionModel model from: {unet_model_dir}")
model = UNet2DConditionModel.from_pretrained(unet_model_dir, **model_opts).to(self.device)
model = optimizer.optimize_checkpoint(model, torch_inference)
return model
def get_input_names(self):
return ["sample", "timestep", "encoder_hidden_states", "text_embeds", "time_ids"]
def get_output_names(self):
return ["latent"]
def get_dynamic_axes(self):
xB = "2B" if self.xB == 2 else "B"
return {
"sample": {0: xB, 2: "H", 3: "W"},
"encoder_hidden_states": {0: xB},
"latent": {0: xB, 2: "H", 3: "W"},
"text_embeds": {0: xB},
"time_ids": {0: xB},
}
def get_input_profile(self, batch_size, image_height, image_width, static_batch, static_shape):
# WAR to enable inference for H/W that are not multiples of 16
# If building with Dynamic Shapes: ensure image height and width are not multiples of 16 for ONNX export and TensorRT engine build
if not static_shape:
image_height = image_height - 8 if image_height % 16 == 0 else image_height
image_width = image_width - 8 if image_width % 16 == 0 else image_width
latent_height, latent_width = self.check_dims(batch_size, image_height, image_width)
min_batch, max_batch, _, _, _, _, min_latent_height, max_latent_height, min_latent_width, max_latent_width = (
self.get_minmax_dims(batch_size, image_height, image_width, static_batch, static_shape)
)
return {
"sample": [
(self.xB * min_batch, self.unet_dim, min_latent_height, min_latent_width),
(self.xB * batch_size, self.unet_dim, latent_height, latent_width),
(self.xB * max_batch, self.unet_dim, max_latent_height, max_latent_width),
],
"encoder_hidden_states": [
(self.xB * min_batch, self.text_maxlen, self.embedding_dim),
(self.xB * batch_size, self.text_maxlen, self.embedding_dim),
(self.xB * max_batch, self.text_maxlen, self.embedding_dim),
],
"text_embeds": [(self.xB * min_batch, 1280), (self.xB * batch_size, 1280), (self.xB * max_batch, 1280)],
"time_ids": [
(self.xB * min_batch, self.time_dim),
(self.xB * batch_size, self.time_dim),
(self.xB * max_batch, self.time_dim),
],
}
def get_shape_dict(self, batch_size, image_height, image_width):
latent_height, latent_width = self.check_dims(batch_size, image_height, image_width)
return {
"sample": (self.xB * batch_size, self.unet_dim, latent_height, latent_width),
"encoder_hidden_states": (self.xB * batch_size, self.text_maxlen, self.embedding_dim),
"latent": (self.xB * batch_size, 4, latent_height, latent_width),
"text_embeds": (self.xB * batch_size, 1280),
"time_ids": (self.xB * batch_size, self.time_dim),
}
def get_sample_input(self, batch_size, image_height, image_width, static_shape):
# WAR to enable inference for H/W that are not multiples of 16
# If building with Dynamic Shapes: ensure image height and width are not multiples of 16 for ONNX export and TensorRT engine build
if not static_shape:
image_height = image_height - 8 if image_height % 16 == 0 else image_height
image_width = image_width - 8 if image_width % 16 == 0 else image_width
latent_height, latent_width = self.check_dims(batch_size, image_height, image_width)
dtype = torch.float16 if self.fp16 else torch.float32
return (
torch.randn(
self.xB * batch_size, self.unet_dim, latent_height, latent_width, dtype=dtype, device=self.device
),
torch.tensor([1.0], dtype=dtype, device=self.device),
torch.randn(self.xB * batch_size, self.text_maxlen, self.embedding_dim, dtype=dtype, device=self.device),
{
"added_cond_kwargs": {
"text_embeds": torch.randn(self.xB * batch_size, 1280, dtype=dtype, device=self.device),
"time_ids": torch.randn(self.xB * batch_size, self.time_dim, dtype=dtype, device=self.device),
}
},
)
def optimize(self, onnx_graph):
if self.fp8:
return super().optimize(onnx_graph, modify_fp8_graph=True)
if self.int8:
return super().optimize(onnx_graph, fuse_mha_qkv_int8=True)
return super().optimize(onnx_graph)
class UNetXLModelControlNet(UNetXLModel):
def __init__(
self,
version,
pipeline,
device,
hf_token,
verbose,
framework_model_dir,
fp16=False,
int8=False,
fp8=False,
max_batch_size=16,
text_maxlen=77,
controlnets=None,
do_classifier_free_guidance=False,
):
super().__init__(
version=version,
pipeline=pipeline,
device=device,
hf_token=hf_token,
verbose=verbose,
framework_model_dir=framework_model_dir,
fp16=fp16,
int8=int8,
fp8=fp8,
max_batch_size=max_batch_size,
text_maxlen=text_maxlen,
do_classifier_free_guidance=do_classifier_free_guidance,
)
self.controlnets = load.get_path(version, pipeline, controlnets) if controlnets else None
def get_model(self, torch_inference=""):
model_opts = {"variant": "fp16", "torch_dtype": torch.float16} if self.fp16 else {}
unet_model = UNet2DConditionModel.from_pretrained(
self.path,
subfolder=self.subfolder,
use_safetensors=self.hf_safetensor,
token=self.hf_token,
**model_opts,
).to(self.device)
cnet_model_opts = {"torch_dtype": torch.float16} if self.fp16 else {}
controlnets = torch.nn.ModuleList(
[ControlNetModel.from_pretrained(path, **cnet_model_opts).to(self.device) for path in self.controlnets]
)
# FIXME - cache UNet2DConditionControlNetModel
model = UNet2DConditionControlNetModel(unet_model, controlnets)
model = optimizer.optimize_checkpoint(model, torch_inference)
return model
def get_input_names(self):
return ["sample", "timestep", "encoder_hidden_states", "images", "controlnet_scales", "text_embeds", "time_ids"]
def get_dynamic_axes(self):
xB = "2B" if self.xB == 2 else "B"
result = super().get_dynamic_axes()
result["images"] = {1: xB, 3: "8H", 4: "8W"}
return result
def get_input_profile(self, batch_size, image_height, image_width, static_batch, static_shape):
min_batch, max_batch, min_image_height, max_image_height, min_image_width, max_image_width, _, _, _, _ = (
self.get_minmax_dims(batch_size, image_height, image_width, static_batch, static_shape)
)
result = super().get_input_profile(batch_size, image_height, image_width, static_batch, static_shape)
result["images"] = [
(len(self.controlnets), self.xB * min_batch, 3, min_image_height, min_image_width),
(len(self.controlnets), self.xB * batch_size, 3, image_height, image_width),
(len(self.controlnets), self.xB * max_batch, 3, max_image_height, max_image_width),
]
return result
def get_shape_dict(self, batch_size, image_height, image_width):
result = super().get_shape_dict(batch_size, image_height, image_width)
result["images"] = (len(self.controlnets), self.xB * batch_size, 3, image_height, image_width)
return result
def get_sample_input(self, batch_size, image_height, image_width, static_shape):
dtype = torch.float16 if self.fp16 else torch.float32
result = super().get_sample_input(batch_size, image_height, image_width, static_shape)
result = (
result[:-1]
+ (
torch.randn(
len(self.controlnets),
self.xB * batch_size,
3,
image_height,
image_width,
dtype=dtype,
device=self.device,
), # images
torch.randn(len(self.controlnets), dtype=dtype, device=self.device), # controlnet_scales
)
+ result[-1:]
)
return result
class UNetTemporalModel(base_model.BaseModel):
def __init__(
self,
version,
pipeline,
device,
hf_token,
verbose,
framework_model_dir,
fp16=False,
fp8=False,
max_batch_size=16,
num_frames=14,
do_classifier_free_guidance=True,
):
super(UNetTemporalModel, self).__init__(
version,
pipeline,
device=device,
hf_token=hf_token,
verbose=verbose,
framework_model_dir=framework_model_dir,
fp16=fp16,
fp8=fp8,
max_batch_size=max_batch_size,
embedding_dim=get_unet_embedding_dim(version, pipeline),
)
self.subfolder = "unet"
self.unet_dim = 4
self.num_frames = num_frames
self.out_channels = 4
self.cross_attention_dim = 1024
self.xB = 2 if do_classifier_free_guidance else 1 # batch multiplier
def get_model(self, torch_inference=""):
model_opts = {"torch_dtype": torch.float16} if self.fp16 else {}
unet_model_dir = load.get_checkpoint_dir(self.framework_model_dir, self.version, self.pipeline, self.subfolder)
if not load.is_model_cached(unet_model_dir, model_opts, self.hf_safetensor):
model = UNetSpatioTemporalConditionModel.from_pretrained(
self.path,
subfolder=self.subfolder,
use_safetensors=self.hf_safetensor,
token=self.hf_token,
**model_opts,
).to(self.device)
model.save_pretrained(unet_model_dir, **model_opts)
else:
print(f"[I] Load UNetSpatioTemporalConditionModel model from: {unet_model_dir}")
model = UNetSpatioTemporalConditionModel.from_pretrained(unet_model_dir, **model_opts).to(self.device)
model = optimizer.optimize_checkpoint(model, torch_inference)
return model
def get_input_names(self):
return ["sample", "timestep", "encoder_hidden_states", "added_time_ids"]
def get_output_names(self):
return ["latent"]
def get_dynamic_axes(self):
xB = str(self.xB) + "B"
return {
"sample": {0: xB, 1: "num_frames", 3: "H", 4: "W"},
"encoder_hidden_states": {0: xB},
"added_time_ids": {0: xB},
}
def get_input_profile(self, batch_size, image_height, image_width, static_batch, static_shape):
latent_height, latent_width = self.check_dims(batch_size, image_height, image_width)
(
min_batch,
max_batch,
min_image_height,
max_image_height,
min_image_width,
max_image_width,
min_latent_height,
max_latent_height,
min_latent_width,
max_latent_width,
) = self.get_minmax_dims(batch_size, image_height, image_width, static_batch, static_shape)
return {
"sample": [
(self.xB * min_batch, self.num_frames, 2 * self.out_channels, min_latent_height, min_latent_width),
(self.xB * batch_size, self.num_frames, 2 * self.out_channels, latent_height, latent_width),
(self.xB * max_batch, self.num_frames, 2 * self.out_channels, max_latent_height, max_latent_width),
],
"encoder_hidden_states": [
(self.xB * min_batch, 1, self.cross_attention_dim),
(self.xB * batch_size, 1, self.cross_attention_dim),
(self.xB * max_batch, 1, self.cross_attention_dim),
],
"added_time_ids": [(self.xB * min_batch, 3), (self.xB * batch_size, 3), (self.xB * max_batch, 3)],
}
def get_shape_dict(self, batch_size, image_height, image_width):
latent_height, latent_width = self.check_dims(batch_size, image_height, image_width)
return {
"sample": (self.xB * batch_size, self.num_frames, 2 * self.out_channels, latent_height, latent_width),
"timestep": (1,),
"encoder_hidden_states": (self.xB * batch_size, 1, self.cross_attention_dim),
"added_time_ids": (self.xB * batch_size, 3),
}
def get_sample_input(self, batch_size, image_height, image_width, static_shape):
# TODO chunk_size if forward_chunking is used
latent_height, latent_width = self.check_dims(batch_size, image_height, image_width)
dtype = torch.float16 if self.fp16 else torch.float32
return (
torch.randn(
self.xB * batch_size,
self.num_frames,
2 * self.out_channels,
latent_height,
latent_width,
dtype=dtype,
device=self.device,
),
torch.tensor([1.0], dtype=torch.float32, device=self.device),
torch.randn(self.xB * batch_size, 1, self.cross_attention_dim, dtype=dtype, device=self.device),
torch.randn(self.xB * batch_size, 3, dtype=dtype, device=self.device),
)
def optimize(self, onnx_graph):
return super().optimize(onnx_graph, modify_fp8_graph=self.fp8)
class UNetCascadeModel(base_model.BaseModel):
def __init__(
self,
version,
pipeline,
device,
hf_token,
verbose,
framework_model_dir,
fp16=False,
bf16=False,
max_batch_size=16,
text_maxlen=77,
do_classifier_free_guidance=False,
compression_factor=42,
latent_dim_scale=10.67,
image_embedding_dim=768,
lite=False,
):
super(UNetCascadeModel, self).__init__(
version,
pipeline,
device=device,
hf_token=hf_token,
verbose=verbose,
framework_model_dir=framework_model_dir,
fp16=fp16,
bf16=bf16,
max_batch_size=max_batch_size,
text_maxlen=text_maxlen,
embedding_dim=get_unet_embedding_dim(version, pipeline),
compression_factor=compression_factor,
)
self.is_prior = True if pipeline.is_cascade_prior() else False
self.subfolder = "prior" if self.is_prior else "decoder"
if lite:
self.subfolder += "_lite"
self.prior_dim = 16
self.decoder_dim = 4
self.xB = 2 if do_classifier_free_guidance else 1 # batch multiplier
self.latent_dim_scale = latent_dim_scale
self.min_latent_shape = self.min_image_shape // self.compression_factor
self.max_latent_shape = self.max_image_shape // self.compression_factor
self.do_constant_folding = False
self.image_embedding_dim = image_embedding_dim
def get_model(self, torch_inference=""):
# FP16 variant doesn't exist
model_opts = {"torch_dtype": torch.float16} if self.fp16 else {}
model_opts = {"variant": "bf16", "torch_dtype": torch.bfloat16} if self.bf16 else model_opts
unet_model_dir = load.get_checkpoint_dir(self.framework_model_dir, self.version, self.pipeline, self.subfolder)
if not load.is_model_cached(unet_model_dir, model_opts, self.hf_safetensor):
model = StableCascadeUNet.from_pretrained(
self.path,
subfolder=self.subfolder,
use_safetensors=self.hf_safetensor,
token=self.hf_token,
**model_opts,
).to(self.device)
model.save_pretrained(unet_model_dir, **model_opts)
else:
print(f"[I] Load Stable Cascade UNet pytorch model from: {unet_model_dir}")
model = StableCascadeUNet.from_pretrained(unet_model_dir, **model_opts).to(self.device)
model = optimizer.optimize_checkpoint(model, torch_inference)
return model
def get_input_names(self):
if self.is_prior:
return ["sample", "timestep_ratio", "clip_text_pooled", "clip_text", "clip_img"]
else:
return ["sample", "timestep_ratio", "clip_text_pooled", "effnet"]
def get_output_names(self):
return ["latent"]
def get_dynamic_axes(self):
xB = "2B" if self.xB == 2 else "B"
if self.is_prior:
return {
"sample": {0: xB, 2: "H", 3: "W"},
"timestep_ratio": {0: xB},
"clip_text_pooled": {0: xB},
"clip_text": {0: xB},
"clip_img": {0: xB},
"latent": {0: xB, 2: "H", 3: "W"},
}
else:
return {
"sample": {0: xB, 2: "H", 3: "W"},
"timestep_ratio": {0: xB},
"clip_text_pooled": {0: xB},
"effnet": {0: xB, 2: "H_effnet", 3: "W_effnet"},
"latent": {0: xB, 2: "H", 3: "W"},
}
def get_input_profile(self, batch_size, image_height, image_width, static_batch, static_shape):
latent_height, latent_width = self.check_dims(batch_size, image_height, image_width)
min_batch, max_batch, _, _, _, _, min_latent_height, max_latent_height, min_latent_width, max_latent_width = (
self.get_minmax_dims(batch_size, image_height, image_width, static_batch, static_shape)
)
if self.is_prior:
return {
"sample": [
(self.xB * min_batch, self.prior_dim, min_latent_height, min_latent_width),
(self.xB * batch_size, self.prior_dim, latent_height, latent_width),
(self.xB * max_batch, self.prior_dim, max_latent_height, max_latent_width),
],
"timestep_ratio": [(self.xB * min_batch,), (self.xB * batch_size,), (self.xB * max_batch,)],
"clip_text_pooled": [
(self.xB * min_batch, 1, self.embedding_dim),
(self.xB * batch_size, 1, self.embedding_dim),
(self.xB * max_batch, 1, self.embedding_dim),
],
"clip_text": [
(self.xB * min_batch, self.text_maxlen, self.embedding_dim),
(self.xB * batch_size, self.text_maxlen, self.embedding_dim),
(self.xB * max_batch, self.text_maxlen, self.embedding_dim),
],
"clip_img": [
(self.xB * min_batch, 1, self.image_embedding_dim),
(self.xB * batch_size, 1, self.image_embedding_dim),
(self.xB * max_batch, 1, self.image_embedding_dim),
],
}
else:
return {
"sample": [
(
self.xB * min_batch,
self.decoder_dim,
int(min_latent_height * self.latent_dim_scale),
int(min_latent_width * self.latent_dim_scale),
),
(
self.xB * batch_size,
self.decoder_dim,
int(latent_height * self.latent_dim_scale),
int(latent_width * self.latent_dim_scale),
),
(
self.xB * max_batch,
self.decoder_dim,
int(max_latent_height * self.latent_dim_scale),
int(max_latent_width * self.latent_dim_scale),
),
],
"timestep_ratio": [(self.xB * min_batch,), (self.xB * batch_size,), (self.xB * max_batch,)],
"clip_text_pooled": [
(self.xB * min_batch, 1, self.embedding_dim),
(self.xB * batch_size, 1, self.embedding_dim),
(self.xB * max_batch, 1, self.embedding_dim),
],
"effnet": [
(self.xB * min_batch, self.prior_dim, min_latent_height, min_latent_width),
(self.xB * batch_size, self.prior_dim, latent_height, latent_width),
(self.xB * max_batch, self.prior_dim, max_latent_height, max_latent_width),
],
}
def get_shape_dict(self, batch_size, image_height, image_width):
latent_height, latent_width = self.check_dims(batch_size, image_height, image_width)
if self.is_prior:
return {
"sample": (self.xB * batch_size, self.prior_dim, latent_height, latent_width),
"timestep_ratio": (self.xB * batch_size,),
"clip_text_pooled": (self.xB * batch_size, 1, self.embedding_dim),
"clip_text": (self.xB * batch_size, self.text_maxlen, self.embedding_dim),
"clip_img": (self.xB * batch_size, 1, self.image_embedding_dim),
"latent": (self.xB * batch_size, self.prior_dim, latent_height, latent_width),
}
else:
return {
"sample": (
self.xB * batch_size,
self.decoder_dim,
int(latent_height * self.latent_dim_scale),
int(latent_width * self.latent_dim_scale),
),
"timestep_ratio": (self.xB * batch_size,),
"clip_text_pooled": (self.xB * batch_size, 1, self.embedding_dim),
"effnet": (self.xB * batch_size, self.prior_dim, latent_height, latent_width),
"latent": (
self.xB * batch_size,
self.decoder_dim,
int(latent_height * self.latent_dim_scale),
int(latent_width * self.latent_dim_scale),
),
}
def get_sample_input(self, batch_size, image_height, image_width, static_shape):
latent_height, latent_width = self.check_dims(batch_size, image_height, image_width)
dtype = torch.float16 if self.fp16 else torch.bfloat16 if self.bf16 else torch.float32
if self.is_prior:
return (
torch.randn(batch_size, self.prior_dim, latent_height, latent_width, dtype=dtype, device=self.device),
torch.tensor([1.0] * batch_size, dtype=dtype, device=self.device),
torch.randn(batch_size, 1, self.embedding_dim, dtype=dtype, device=self.device),
{
"clip_text": torch.randn(
batch_size, self.text_maxlen, self.embedding_dim, dtype=dtype, device=self.device
),
"clip_img": torch.randn(batch_size, 1, self.image_embedding_dim, dtype=dtype, device=self.device),
},
)
else:
return (
torch.randn(
batch_size,
self.decoder_dim,
int(latent_height * self.latent_dim_scale),
int(latent_width * self.latent_dim_scale),
dtype=dtype,
device=self.device,
),
torch.tensor([1.0] * batch_size, dtype=dtype, device=self.device),
torch.randn(batch_size, 1, self.embedding_dim, dtype=dtype, device=self.device),
{
"effnet": torch.randn(
batch_size, self.prior_dim, latent_height, latent_width, dtype=dtype, device=self.device
),
},
)
demo/Diffusion/demo_diffusion/model/vae.py
+473
View File
@@ -0,0 +1,473 @@
#
# SPDX-FileCopyrightText: Copyright (c) 1993-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
import os
import torch
from huggingface_hub import hf_hub_download
from safetensors import safe_open
from demo_diffusion.dynamic_import import import_from_diffusers
from demo_diffusion.model import base_model, load, optimizer
from demo_diffusion.utils_sd3.other_impls import load_into
from demo_diffusion.utils_sd3.sd3_impls import SDVAE
# List of models to import from diffusers.models
models_to_import = [
"AutoencoderKL",
"AutoencoderKLTemporalDecoder",
]
for model in models_to_import:
globals()[model] = import_from_diffusers(model, "diffusers.models")
class VAEModel(base_model.BaseModel):
def __init__(
self,
version,
pipeline,
device,
hf_token,
verbose,
framework_model_dir,
fp16=False,
tf32=False,
bf16=False,
max_batch_size=16,
):
super(VAEModel, self).__init__(
version,
pipeline,
device=device,
hf_token=hf_token,
verbose=verbose,
framework_model_dir=framework_model_dir,
fp16=fp16,
tf32=tf32,
bf16=bf16,
max_batch_size=max_batch_size,
)
self.subfolder = "vae"
self.vae_decoder_model_dir = load.get_checkpoint_dir(
self.framework_model_dir, self.version, self.pipeline, self.subfolder
)
if not os.path.exists(self.vae_decoder_model_dir):
self.config = AutoencoderKL.load_config(self.path, subfolder=self.subfolder, token=self.hf_token)
else:
print(f"[I] Load AutoencoderKL (decoder) config from: {self.vae_decoder_model_dir}")
self.config = AutoencoderKL.load_config(self.vae_decoder_model_dir)
def get_model(self, torch_inference=""):
model_opts = (
{"torch_dtype": torch.float16} if self.fp16 else {"torch_dtype": torch.bfloat16} if self.bf16 else {}
)
if not load.is_model_cached(self.vae_decoder_model_dir, model_opts, self.hf_safetensor):
model = AutoencoderKL.from_pretrained(
self.path,
subfolder=self.subfolder,
use_safetensors=self.hf_safetensor,
token=self.hf_token,
**model_opts,
).to(self.device)
model.save_pretrained(self.vae_decoder_model_dir, **model_opts)
else:
print(f"[I] Load AutoencoderKL (decoder) model from: {self.vae_decoder_model_dir}")
model = AutoencoderKL.from_pretrained(self.vae_decoder_model_dir, **model_opts).to(self.device)
model.forward = model.decode
model = optimizer.optimize_checkpoint(model, torch_inference)
return model
def get_input_names(self):
return ["latent"]
def get_output_names(self):
return ["images"]
def get_dynamic_axes(self):
return {"latent": {0: "B", 2: "H", 3: "W"}, "images": {0: "B", 2: "8H", 3: "8W"}}
def get_input_profile(self, batch_size, image_height, image_width, static_batch, static_shape):
latent_height, latent_width = self.check_dims(batch_size, image_height, image_width)
min_batch, max_batch, _, _, _, _, min_latent_height, max_latent_height, min_latent_width, max_latent_width = (
self.get_minmax_dims(batch_size, image_height, image_width, static_batch, static_shape)
)
return {
"latent": [
(min_batch, self.config["latent_channels"], min_latent_height, min_latent_width),
(batch_size, self.config["latent_channels"], latent_height, latent_width),
(max_batch, self.config["latent_channels"], max_latent_height, max_latent_width),
]
}
def get_shape_dict(self, batch_size, image_height, image_width):
latent_height, latent_width = self.check_dims(batch_size, image_height, image_width)
return {
"latent": (batch_size, self.config["latent_channels"], latent_height, latent_width),
"images": (batch_size, 3, image_height, image_width),
}
def get_sample_input(self, batch_size, image_height, image_width, static_shape):
latent_height, latent_width = self.check_dims(batch_size, image_height, image_width)
dtype = torch.float16 if self.fp16 else torch.bfloat16 if self.bf16 else torch.float32
return torch.randn(
batch_size, self.config["latent_channels"], latent_height, latent_width, dtype=dtype, device=self.device
)
class SD3_VAEDecoderModel(base_model.BaseModel):
def __init__(
self,
version,
pipeline,
device,
hf_token,
verbose,
framework_model_dir,
max_batch_size,
fp16=False,
):
super(SD3_VAEDecoderModel, self).__init__(
version,
pipeline,
device=device,
hf_token=hf_token,
verbose=verbose,
framework_model_dir=framework_model_dir,
fp16=fp16,
max_batch_size=max_batch_size,
)
self.subfolder = "sd3"
def get_model(self, torch_inference=""):
dtype = torch.float16 if self.fp16 else torch.float32
sd3_model_dir = load.get_checkpoint_dir(self.framework_model_dir, self.version, self.pipeline, self.subfolder)
sd3_filename = "sd3_medium.safetensors"
sd3_model_path = f"{sd3_model_dir}/{sd3_filename}"
if not os.path.exists(sd3_model_path):
hf_hub_download(repo_id=self.path, filename=sd3_filename, local_dir=sd3_model_dir)
with safe_open(sd3_model_path, framework="pt", device=self.device) as f:
model = SDVAE(device=self.device, dtype=dtype).eval().cuda()
prefix = ""
if any(k.startswith("first_stage_model.") for k in f.keys()):
prefix = "first_stage_model."
load_into(f, model, prefix, self.device, dtype)
model.forward = model.decode
model = optimizer.optimize_checkpoint(model, torch_inference)
return model
def get_input_names(self):
return ["latent"]
def get_output_names(self):
return ["images"]
def get_dynamic_axes(self):
return {"latent": {0: "B", 2: "H", 3: "W"}, "images": {0: "B", 2: "8H", 3: "8W"}}
def get_input_profile(self, batch_size, image_height, image_width, static_batch, static_shape):
latent_height, latent_width = self.check_dims(batch_size, image_height, image_width)
min_batch, max_batch, _, _, _, _, min_latent_height, max_latent_height, min_latent_width, max_latent_width = (
self.get_minmax_dims(batch_size, image_height, image_width, static_batch, static_shape)
)
return {
"latent": [
(min_batch, 16, min_latent_height, min_latent_width),
(batch_size, 16, latent_height, latent_width),
(max_batch, 16, max_latent_height, max_latent_width),
]
}
def get_shape_dict(self, batch_size, image_height, image_width):
latent_height, latent_width = self.check_dims(batch_size, image_height, image_width)
return {
"latent": (batch_size, 16, latent_height, latent_width),
"images": (batch_size, 3, image_height, image_width),
}
def get_sample_input(self, batch_size, image_height, image_width, static_shape):
latent_height, latent_width = self.check_dims(batch_size, image_height, image_width)
dtype = torch.float16 if self.fp16 else torch.float32
return torch.randn(batch_size, 16, latent_height, latent_width, dtype=dtype, device=self.device)
class VAEDecTemporalModel(base_model.BaseModel):
def __init__(
self,
version,
pipeline,
device,
hf_token,
verbose,
framework_model_dir,
max_batch_size=16,
decode_chunk_size=14,
):
super(VAEDecTemporalModel, self).__init__(
version,
pipeline,
device=device,
hf_token=hf_token,
verbose=verbose,
framework_model_dir=framework_model_dir,
max_batch_size=max_batch_size,
)
self.subfolder = "vae"
self.decode_chunk_size = decode_chunk_size
def get_model(self, torch_inference=""):
vae_decoder_model_path = load.get_checkpoint_dir(
self.framework_model_dir, self.version, self.pipeline, self.subfolder
)
if not os.path.exists(vae_decoder_model_path):
model = AutoencoderKLTemporalDecoder.from_pretrained(
self.path, subfolder=self.subfolder, use_safetensors=self.hf_safetensor, token=self.hf_token
).to(self.device)
model.save_pretrained(vae_decoder_model_path)
else:
print(f"[I] Load AutoencoderKLTemporalDecoder model from: {vae_decoder_model_path}")
model = AutoencoderKLTemporalDecoder.from_pretrained(vae_decoder_model_path).to(self.device)
model.forward = model.decode
model = optimizer.optimize_checkpoint(model, torch_inference)
return model
def get_input_names(self):
return ["latent", "num_frames_in"]
def get_output_names(self):
return ["frames"]
def get_dynamic_axes(self):
return {"latent": {0: "num_frames_in", 2: "H", 3: "W"}, "frames": {0: "num_frames_in", 2: "8H", 3: "8W"}}
def get_input_profile(self, batch_size, image_height, image_width, static_batch, static_shape):
latent_height, latent_width = self.check_dims(batch_size, image_height, image_width)
assert batch_size == 1
_, _, _, _, _, _, min_latent_height, max_latent_height, min_latent_width, max_latent_width = (
self.get_minmax_dims(batch_size, image_height, image_width, static_batch, static_shape)
)
return {
"latent": [
(1, 4, min_latent_height, min_latent_width),
(self.decode_chunk_size, 4, latent_height, latent_width),
(self.decode_chunk_size, 4, max_latent_height, max_latent_width),
],
"num_frames_in": [(1,), (1,), (1,)],
}
def get_shape_dict(self, batch_size, image_height, image_width):
latent_height, latent_width = self.check_dims(batch_size, image_height, image_width)
assert batch_size == 1
return {
"latent": (self.decode_chunk_size, 4, latent_height, latent_width),
#'num_frames_in': (1,),
"frames": (self.decode_chunk_size, 3, image_height, image_width),
}
def get_sample_input(self, batch_size, image_height, image_width):
latent_height, latent_width = self.check_dims(batch_size, image_height, image_width)
assert batch_size == 1
return (
torch.randn(
self.decode_chunk_size, 4, latent_height, latent_width, dtype=torch.float32, device=self.device
),
self.decode_chunk_size,
)
class TorchVAEEncoder(torch.nn.Module):
def __init__(
self,
version,
pipeline,
hf_token,
device,
path,
framework_model_dir,
subfolder,
fp16=False,
bf16=False,
hf_safetensor=False,
):
super().__init__()
model_opts = {"torch_dtype": torch.float16} if fp16 else {"torch_dtype": torch.bfloat16} if bf16 else {}
vae_encoder_model_dir = load.get_checkpoint_dir(framework_model_dir, version, pipeline, subfolder)
if not load.is_model_cached(vae_encoder_model_dir, model_opts, hf_safetensor):
self.vae_encoder = AutoencoderKL.from_pretrained(
path, subfolder="vae", use_safetensors=hf_safetensor, token=hf_token, **model_opts
).to(device)
self.vae_encoder.save_pretrained(vae_encoder_model_dir, **model_opts)
else:
print(f"[I] Load AutoencoderKL (encoder) model from: {vae_encoder_model_dir}")
self.vae_encoder = AutoencoderKL.from_pretrained(vae_encoder_model_dir, **model_opts).to(device)
def forward(self, x):
return self.vae_encoder.encode(x).latent_dist.sample()
class VAEEncoderModel(base_model.BaseModel):
def __init__(
self,
version,
pipeline,
device,
hf_token,
verbose,
framework_model_dir,
fp16=False,
tf32=False,
bf16=False,
max_batch_size=16,
):
super(VAEEncoderModel, self).__init__(
version,
pipeline,
device=device,
hf_token=hf_token,
verbose=verbose,
framework_model_dir=framework_model_dir,
fp16=fp16,
tf32=tf32,
bf16=bf16,
max_batch_size=max_batch_size,
)
self.subfolder = "vae"
self.vae_encoder_model_dir = load.get_checkpoint_dir(
framework_model_dir, version, self.pipeline, self.subfolder
)
if not os.path.exists(self.vae_encoder_model_dir):
self.config = AutoencoderKL.load_config(self.path, subfolder=self.subfolder, token=self.hf_token)
else:
print(f"[I] Load AutoencoderKL (encoder) config from: {self.vae_encoder_model_dir}")
self.config = AutoencoderKL.load_config(self.vae_encoder_model_dir)
def get_model(self, torch_inference=""):
vae_encoder = TorchVAEEncoder(
self.version,
self.pipeline,
self.hf_token,
self.device,
self.path,
self.framework_model_dir,
self.subfolder,
self.fp16,
self.bf16,
hf_safetensor=self.hf_safetensor,
)
return vae_encoder
def get_input_names(self):
return ["images"]
def get_output_names(self):
return ["latent"]
def get_dynamic_axes(self):
return {"images": {0: "B", 2: "8H", 3: "8W"}, "latent": {0: "B", 2: "H", 3: "W"}}
def get_input_profile(self, batch_size, image_height, image_width, static_batch, static_shape):
assert batch_size >= self.min_batch and batch_size <= self.max_batch
min_batch = batch_size if static_batch else self.min_batch
max_batch = batch_size if static_batch else self.max_batch
self.check_dims(batch_size, image_height, image_width)
min_batch, max_batch, min_image_height, max_image_height, min_image_width, max_image_width, _, _, _, _ = (
self.get_minmax_dims(batch_size, image_height, image_width, static_batch, static_shape)
)
return {
"images": [
(min_batch, 3, min_image_height, min_image_width),
(batch_size, 3, image_height, image_width),
(max_batch, 3, max_image_height, max_image_width),
],
}
def get_shape_dict(self, batch_size, image_height, image_width):
latent_height, latent_width = self.check_dims(batch_size, image_height, image_width)
return {
"images": (batch_size, 3, image_height, image_width),
"latent": (batch_size, self.config["latent_channels"], latent_height, latent_width),
}
def get_sample_input(self, batch_size, image_height, image_width, static_shape):
self.check_dims(batch_size, image_height, image_width)
dtype = torch.float16 if self.fp16 else torch.bfloat16 if self.bf16 else torch.float32
return torch.randn(batch_size, 3, image_height, image_width, dtype=dtype, device=self.device)
class SD3_VAEEncoderModel(VAEEncoderModel):
def __init__(
self,
version,
pipeline,
device,
hf_token,
verbose,
framework_model_dir,
max_batch_size,
fp16=False,
):
super(SD3_VAEEncoderModel, self).__init__(
version,
pipeline,
device=device,
hf_token=hf_token,
verbose=verbose,
framework_model_dir=framework_model_dir,
fp16=fp16,
max_batch_size=max_batch_size,
)
self.subfolder = "sd3"
def get_model(self, torch_inference=""):
dtype = torch.float16 if self.fp16 else torch.float32
sd3_model_dir = load.get_checkpoint_dir(self.framework_model_dir, self.version, self.pipeline, self.subfolder)
sd3_filename = "sd3_medium.safetensors"
sd3_model_path = f"{sd3_model_dir}/{sd3_filename}"
if not os.path.exists(sd3_model_path):
hf_hub_download(repo_id=self.path, filename=sd3_filename, local_dir=sd3_model_dir)
with safe_open(sd3_model_path, framework="pt", device=self.device) as f:
model = SDVAE(device=self.device, dtype=dtype).eval().cuda()
prefix = ""
if any(k.startswith("first_stage_model.") for k in f.keys()):
prefix = "first_stage_model."
load_into(f, model, prefix, self.device, dtype)
model.forward = model.encode
model = optimizer.optimize_checkpoint(model, torch_inference)
return model
def get_input_profile(self, batch_size, image_height, image_width, static_batch, static_shape):
min_batch, max_batch, _, _, _, _, _, _, _, _ = self.get_minmax_dims(
batch_size, image_height, image_width, static_batch, static_shape
)
return {
"images": [
(min_batch, 3, image_height, image_width),
(batch_size, 3, image_height, image_width),
(max_batch, 3, image_height, image_width),
]
}
def get_shape_dict(self, batch_size, image_height, image_width):
latent_height, latent_width = self.check_dims(batch_size, image_height, image_width)
return {
"images": (batch_size, 3, image_height, image_width),
"latent": (batch_size, 16, latent_height, latent_width),
}
def get_sample_input(self, batch_size, image_height, image_width, static_shape):
dtype = torch.float16 if self.fp16 else torch.float32
return torch.randn(batch_size, 3, image_height, image_width, dtype=dtype, device=self.device)
demo/Diffusion/demo_diffusion/path/__init__.py
+21
View File
@@ -0,0 +1,21 @@
#
# SPDX-FileCopyrightText: Copyright (c) 1993-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
from demo_diffusion.path.dd_path import DDPath
from demo_diffusion.path.resolve_path import resolve_path
__all__ = ["DDPath", "resolve_path"]
demo/Diffusion/demo_diffusion/path/dd_path.py
+54
View File
@@ -0,0 +1,54 @@
#
# SPDX-FileCopyrightText: Copyright (c) 1993-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
"""
Define a data structure for storing various paths used in DemoDiffusion.
"""
import dataclasses
import os
from typing import Dict
@dataclasses.dataclass
class DDPath:
"""Data class that stores various paths used in DemoDiffusion."""
# Input paths.
# NOTE: Pre-resolved checkpoint paths are currently not used. Pipelines will be updated to using pre-resolved
# checkpoint paths soon in the future.
model_name_to_checkpoint_path: Dict[str, str] = dataclasses.field(default_factory=dict)
model_name_to_optimized_onnx_path: Dict[str, str] = dataclasses.field(default_factory=dict)
model_name_to_engine_path: Dict[str, str] = dataclasses.field(default_factory=dict)
# Artifact paths.
model_name_to_unoptimized_onnx_path: Dict[str, str] = dataclasses.field(default_factory=dict)
model_name_to_weights_map_path: Dict[str, str] = dataclasses.field(default_factory=dict)
model_name_to_refit_weights_path: Dict[str, str] = dataclasses.field(default_factory=dict)
model_name_to_quantized_model_state_dict_path: Dict[str, str] = dataclasses.field(default_factory=dict)
def create_directory(self) -> None:
"""Create directories for all paths, if they do not exist."""
all_paths = [value for name_to_path in dataclasses.astuple(self) for value in name_to_path.values()]
for path in all_paths:
directory = os.path.dirname(path)
# If `path` does not have a directory component, `directory` will be an empty string.
# Only proceed if `directory` is non-empty.
if directory:
os.makedirs(directory, exist_ok=True)
demo/Diffusion/demo_diffusion/path/resolve_path.py
+172
View File
@@ -0,0 +1,172 @@
#
# SPDX-FileCopyrightText: Copyright (c) 1993-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
from __future__ import annotations
import argparse
import hashlib
import os
from typing import Dict, List
import tensorrt as trt
from demo_diffusion import pipeline
from demo_diffusion.path import dd_path
ARTIFACT_CACHE_DIRECTORY = os.path.join(os.getcwd(), "artifacts_cache")
def resolve_path(
model_names: List[str],
args: argparse.Namespace,
pipeline_type: pipeline.PIPELINE_TYPE,
pipeline_uid: str,
) -> dd_path.DDPath:
"""Resolve all paths and store them in a newly constructed dd_path.DDPath object.
Args:
model_names (List[str]): List of model names.
args (argparse.Namespace): Parsed arguments.
Returns:
dd_path.DDPath: Path object containing all the resolved paths.
"""
path = dd_path.DDPath()
model_name_to_model_uri = {
model_name: _resolve_model_uri(model_name, args, pipeline_type, pipeline_uid) for model_name in model_names
}
_resolve_default_path(model_name_to_model_uri, args, path)
_resolve_custom_path(args, path)
path.create_directory()
return path
def _resolve_model_uri(
model_name: str, args: argparse.Namespace, pipeline_type: pipeline.PIPELINE_TYPE, pipeline_uid: str
) -> str:
"""Resolve and return the model URI.
The model URI is a partial path that uniquely identifies the model. It is used to construct various model paths like
artifact cache path, checkpoint path, etc.
"""
# Lora unique ID represents the lora configuration.
if args.lora_path and args.lora_weight:
lora_config_uid = "-".join(
sorted(
[
f"{hashlib.sha256(lora_path.encode()).hexdigest()}-{lora_weight}-{args.lora_scale}"
for lora_path, lora_weight in zip(args.lora_path, args.lora_weight)
if args.lora_path
]
)
)
else:
lora_config_uid = ""
# Quantization config unique ID represents the quantization configuration.
def _is_quantized() -> bool:
"""Return True if model is quantized, False if otherwise.
When quantization flags are set in `args`, only a subset of the models are actually quantized.
"""
is_unet = model_name == "unet"
is_unetxl_base = pipeline_type.is_sd_xl_base() and model_name == "unetxl"
is_flux_transformer = args.version.startswith("flux.1") and model_name == "transformer"
if args.int8:
return is_unet or is_unetxl_base
elif args.fp8:
return is_unet or is_unetxl_base or is_flux_transformer
elif args.fp4:
return is_flux_transformer
else:
return False
if _is_quantized():
if args.int8 or args.fp8:
quantization_config_uid = (
f"{'int8' if args.int8 else 'fp8'}.l{args.quantization_level}.bs2.s{args.denoising_steps}"
f".c{args.calibration_size}.p{args.quantization_percentile}.a{args.quantization_alpha}"
)
else:
quantization_config_uid = "fp4"
else:
quantization_config_uid = ""
# Model unique ID represents the model name and its configuration. It is unique under the same pipeline.
model_uid = "_".join([s for s in [model_name, lora_config_uid, quantization_config_uid] if s])
# Model URI is the concatenation of pipeline unique ID and model unique ID.
model_uri = os.path.join(pipeline_uid, model_uid)
return model_uri
def _resolve_default_path(
model_name_to_model_uri: Dict[str, str], args: argparse.Namespace, path: dd_path.DDPath
) -> None:
"""Resolve the default paths.
Args:
model_name_to_model_uri (Dict[str, str]): Dictionary of model name to model URI.
args (argparse.Namespace): Parsed arguments.
path (dd_path.DDPath): Path object. This object is modified in-place to store all resolved default paths.
"""
for model_name, model_uri in model_name_to_model_uri.items():
# Resolve input paths.
# NOTE: `checkpoint_path` is the only exception where the path points to a directory instead of a file, because
# the HF download-save-load round trip requires a directory.
path.model_name_to_checkpoint_path[model_name] = os.path.join(args.framework_model_dir, model_uri)
path.model_name_to_optimized_onnx_path[model_name] = os.path.join(
args.onnx_dir, model_uri, "model_optimized.onnx"
)
path.model_name_to_engine_path[model_name] = os.path.join(
args.engine_dir, model_uri, f"engine_trt{trt.__version__}.plan"
)
# Resolve artifact paths.
artifact_dir = os.path.join(ARTIFACT_CACHE_DIRECTORY, model_uri)
path.model_name_to_unoptimized_onnx_path[model_name] = os.path.join(artifact_dir, "model_unoptimized.onnx")
path.model_name_to_weights_map_path[model_name] = os.path.join(artifact_dir, "weights_map.json")
path.model_name_to_refit_weights_path[model_name] = os.path.join(artifact_dir, "refit_weights.json")
path.model_name_to_quantized_model_state_dict_path[model_name] = os.path.join(
artifact_dir, "quantized_model_state_dict.json"
)
def _resolve_custom_path(args: argparse.Namespace, path: dd_path.DDPath) -> None:
"""Resolve the custom paths.
If a different path already exists in `path`, it will be overridden.
Args:
args (argparse.Namespace): Parsed arguments.
path (dd_path.DDPath): Path object. This object is modified in-place to store or override all resolved paths.
"""
# Resolve and override custom ONNX paths.
if args.custom_onnx_paths:
for model_name, optimized_onnx_path in args.custom_onnx_paths.items():
path.model_name_to_optimized_onnx_path[model_name] = optimized_onnx_path
# Resolve and override custom engine paths.
if args.custom_engine_paths:
for model_name, engine_path in args.custom_engine_paths.items():
path.model_name_to_engine_path[model_name] = engine_path
demo/Diffusion/demo_diffusion/pipeline/__init__.py
+36
View File
@@ -0,0 +1,36 @@
#
# SPDX-FileCopyrightText: Copyright (c) 1993-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
from demo_diffusion.pipeline.diffusion_pipeline import DiffusionPipeline
from demo_diffusion.pipeline.flux_pipeline import FluxPipeline
from demo_diffusion.pipeline.stable_cascade_pipeline import StableCascadePipeline
from demo_diffusion.pipeline.stable_diffusion_3_pipeline import StableDiffusion3Pipeline
from demo_diffusion.pipeline.stable_diffusion_pipeline import StableDiffusionPipeline
from demo_diffusion.pipeline.stable_video_diffusion_pipeline import (
StableVideoDiffusionPipeline,
)
from demo_diffusion.pipeline.type import PIPELINE_TYPE
__all__ = [
"DiffusionPipeline",
"FluxPipeline",
"StableCascadePipeline",
"StableDiffusion3Pipeline",
"StableDiffusionPipeline",
"StableVideoDiffusionPipeline",
"PIPELINE_TYPE",
]
demo/Diffusion/demo_diffusion/pipeline/calibrate.py
+37
View File
@@ -0,0 +1,37 @@
#
# SPDX-FileCopyrightText: Copyright (c) 1993-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
import os
from diffusers.utils import load_image
def load_calib_prompts(batch_size, calib_data_path):
with open(calib_data_path, "r", encoding="utf-8") as file:
lst = [line.rstrip("\n") for line in file]
return [lst[i : i + batch_size] for i in range(0, len(lst), batch_size)]
def load_calibration_images(folder_path):
images = []
for filename in os.listdir(folder_path):
img_path = os.path.join(folder_path, filename)
if os.path.isfile(img_path):
image = load_image(img_path)
if image is not None:
images.append(image)
return images
demo/Diffusion/diffusion_pipeline.py → demo/Diffusion/demo_diffusion/pipeline/diffusion_pipeline.py
Regular → Executable
+153 -119
View File
@@ -1,6 +1,5 @@
#
# SPDX-FileCopyrightText: Copyright (c) 1993-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-FileCopyrightText: Copyright (c) 1993-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#
# Licensed under the Apache License, Version 2.0 (the "License");
@@ -16,18 +15,20 @@
# limitations under the License.
#
from __future__ import annotations
import abc
import argparse
import gc
import json
import os
import pathlib
from abc import ABC, abstractmethod
from hashlib import md5
from typing import List, Optional
from typing import Any, List, Optional
import modelopt.torch.opt as mto
import modelopt.torch.quantization as mtq
import nvtx
import tensorrt as trt
import torch
from cuda import cudart
from diffusers import (
@@ -42,18 +43,19 @@ from diffusers import (
PNDMScheduler,
UniPCMultistepScheduler,
)
from models import LoraLoader, make_scheduler
from torch.utils.data import DataLoader
from utilities import (
PIPELINE_TYPE,
Engine,
get_refit_weights,
load_calib_prompts,
import demo_diffusion.engine as engine_module
import demo_diffusion.image as image_module
from demo_diffusion.model import (
LoraLoader,
make_scheduler,
merge_loras,
save_image,
unload_model,
unload_torch_model,
)
from utils_modelopt import (
from demo_diffusion.pipeline.calibrate import load_calib_prompts
from demo_diffusion.pipeline.type import PIPELINE_TYPE
from demo_diffusion.utils_modelopt import (
SD_FP8_BF16_FLUX_MMDIT_BMM2_FP8_OUTPUT_CONFIG,
SD_FP8_FP16_DEFAULT_CONFIG,
SD_FP8_FP32_DEFAULT_CONFIG,
@@ -115,6 +117,7 @@ class DiffusionPipeline(ABC):
def __init__(
self,
dd_path,
version="1.5",
pipeline_type=PIPELINE_TYPE.TXT2IMG,
max_batch_size=16,
@@ -140,6 +143,7 @@ class DiffusionPipeline(ABC):
Initializes the Diffusion pipeline.
Args:
dd_path (load_module.DDPath): DDPath object that contains all paths used in DemoDiffusion.
version (str):
The version of the pipeline. Should be one of the values listed in DiffusionPipeline.VALID_DIFFUSION_PIPELINES.
pipeline_type (PIPELINE_TYPE):
@@ -182,6 +186,8 @@ class DiffusionPipeline(ABC):
denoiser_weight_streaming_budget_percentage (`int`, defaults to None):
Weight streaming budget as a percentage of the size of total streamable weights for the denoiser model.
"""
self.dd_path = dd_path
self.denoising_steps = denoising_steps
self.max_batch_size = max_batch_size
@@ -205,6 +211,8 @@ class DiffusionPipeline(ABC):
self.text_encoder_weight_streaming_budget_percentage = text_encoder_weight_streaming_budget_percentage
self.denoiser_weight_streaming_budget_percentage = denoiser_weight_streaming_budget_percentage
self.stages = self.get_model_names(self.pipeline_type)
if not scheduler:
scheduler = 'UniPC' if self.pipeline_type.is_controlnet() else self.SCHEDULER_DEFAULTS.get(version, 'DDIM')
print(f"[I] Autoselected scheduler: {scheduler}")
@@ -262,6 +270,26 @@ class DiffusionPipeline(ABC):
# config to store additional info
self.config = {}
@classmethod
@abc.abstractmethod
def FromArgs(cls, args: argparse.Namespace, pipeline_type: PIPELINE_TYPE) -> DiffusionPipeline:
"""Factory method to construct a concrete pipeline object from parsed arguments."""
raise NotImplementedError("FromArgs cannot be called from the abstract base class.")
@classmethod
@abc.abstractmethod
def get_model_names(cls, pipeline_type: PIPELINE_TYPE) -> List[str]:
"""Return a list of model names used by this pipeline."""
raise NotImplementedError("get_model_names cannot be called from the abstract base class.")
@classmethod
def _get_pipeline_uid(cls, pipeline_type: PIPELINE_TYPE, version: str) -> str:
"""Return the unique ID of this pipeline.
This is typically used to determine the default path for things like engine files, artifacts caches, etc.
"""
return f"{cls.__name__}_{pipeline_type.name}_{version}"
def profile_start(self, name, color='blue'):
if self.nvtx_profile:
self.markers[name] = nvtx.start_range(message=name, color=color)
@@ -294,94 +322,65 @@ class DiffusionPipeline(ABC):
if not self.low_vram:
self.engine[model_name].allocate_buffers(shape_dict=obj.get_shape_dict(batch_size, image_height, image_width), device=self.device)
def _create_directories(self, engine_dir, onnx_dir):
# Create directories if missing
for directory in [engine_dir, onnx_dir]:
if not os.path.exists(directory):
print(f"[I] Create directory: {directory}")
pathlib.Path(directory).mkdir(parents=True)
def _cached_model_name(self, model_name):
if self.pipeline_type.is_inpaint():
model_name += '_inpaint'
return model_name
def _get_onnx_path(self, model_name, onnx_dir, opt=True, suffix=''):
onnx_model_dir = os.path.join(onnx_dir, self._cached_model_name(model_name)+suffix+('.opt' if opt else ''))
os.makedirs(onnx_model_dir, exist_ok=True)
return os.path.join(onnx_model_dir, 'model.onnx')
def _get_engine_path(self, model_name, engine_dir, enable_refit=False, suffix=''):
return os.path.join(engine_dir, self._cached_model_name(model_name)+suffix+('.refit' if enable_refit else '')+'.trt'+trt.__version__+'.plan')
def _get_weights_map_path(self, model_name, onnx_dir):
onnx_model_dir = os.path.join(onnx_dir, self._cached_model_name(model_name)+'.opt')
os.makedirs(onnx_model_dir, exist_ok=True)
return os.path.join(onnx_model_dir, 'weights_map.json')
def _get_refit_nodes_path(self, model_name, onnx_dir, suffix=''):
onnx_model_dir = os.path.join(onnx_dir, self._cached_model_name(model_name)+'.opt')
os.makedirs(onnx_model_dir, exist_ok=True)
return os.path.join(onnx_model_dir, 'refit'+suffix+'.json')
def _get_state_dict_path(self, model_name, onnx_dir, suffix=''):
onnx_model_dir = os.path.join(onnx_dir, self._cached_model_name(model_name)+suffix)
os.makedirs(onnx_model_dir, exist_ok=True)
return os.path.join(onnx_model_dir, 'state_dict.pt')
@abstractmethod
def _initialize_models(self):
def _initialize_models(self, *args, **kwargs):
raise NotImplementedError("Please Implement the _initialize_models method")
def _get_lora_suffix(self):
if self.lora_loader:
return '-' + '-'.join([str(md5(path.encode('utf-8')).hexdigest()) + '-' + ('%.2f' % self.lora_weights[path]) + '-' + ('%.2f' % self.lora_loader.scale) for path in sorted(self.lora_loader.paths)])
return ''
def _prepare_model_configs(self, default_onnx_dir, model_onnx_dirs, engine_dir, enable_refit, int8, fp8, fp4, quantization_level, quantization_percentile, quantization_alpha, calibration_size):
def _prepare_model_configs(
self,
enable_refit,
int8,
fp8,
fp4
):
model_names = self.models.keys()
lora_suffix = self._get_lora_suffix()
self.torch_fallback = dict(zip(model_names, [self.torch_inference or self.config.get(model_name.replace('-','_')+'_torch_fallback', False) for model_name in model_names]))
configs = {}
for model_name in model_names:
if model_onnx_dirs and model_name in model_onnx_dirs:
onnx_dir = model_onnx_dirs[model_name]
print(f"[I] Model {model_name} using onnx dir: {onnx_dir}")
else:
onnx_dir = default_onnx_dir
print(f"[I] Model {model_name} using default onnx dir: {onnx_dir}")
config = {
'do_engine_refit': not self.pipeline_type.is_sd_xl_refiner() and enable_refit and model_name.startswith('unet'),
'do_lora_merge': not enable_refit and self.lora_loader and model_name.startswith('unet'),
'use_int8': False,
'use_fp8': False,
"do_engine_refit": not self.pipeline_type.is_sd_xl_refiner()
and enable_refit
and model_name.startswith("unet"),
"do_lora_merge": not enable_refit and self.lora_loader and model_name.startswith("unet"),
"use_int8": False,
"use_fp8": False,
'use_fp4': False,
}
config['model_suffix'] = lora_suffix if config['do_lora_merge'] else ''
# TODO: Move this to when arguments are first being validated in dd_argparse.py
if int8:
assert self.pipeline_type.is_sd_xl_base() or self.version in ["1.5", "2.1", "2.1-base"], "int8 quantization only supported for SDXL, SD1.5 and SD2.1 pipeline"
if (self.pipeline_type.is_sd_xl() and model_name == 'unetxl') or \
(model_name == 'unet'):
config['use_int8'] = True
config['model_suffix'] += f"-int8.l{quantization_level}.bs2.s{self.denoising_steps}.c{calibration_size}.p{quantization_percentile}.a{quantization_alpha}"
assert self.pipeline_type.is_sd_xl_base() or self.version in [
"1.5",
"2.1",
"2.1-base",
], "int8 quantization only supported for SDXL, SD1.5 and SD2.1 pipeline"
if (self.pipeline_type.is_sd_xl() and model_name == "unetxl") or (model_name == "unet"):
config["use_int8"] = True
elif fp8:
assert self.pipeline_type.is_sd_xl() or self.version in ["1.5", "2.1", "2.1-base"] or self.version.startswith("flux.1"), "fp8 quantization only supported for SDXL, SD1.5, SD2.1 and FLUX pipeline"
if (self.pipeline_type.is_sd_xl() and model_name == 'unetxl') or \
((self.version.startswith("flux.1")) and model_name == 'transformer') or \
(model_name == 'unet'):
config['use_fp8'] = True
config['model_suffix'] += f"-fp8.l{quantization_level}.bs2.s{self.denoising_steps}.c{calibration_size}.p{quantization_percentile}.a{quantization_alpha}"
assert (
self.pipeline_type.is_sd_xl()
or self.version in ["1.5", "2.1", "2.1-base"]
or self.version.startswith("flux.1")
), "fp8 quantization only supported for SDXL, SD1.5, SD2.1 and FLUX pipeline"
if (
(self.pipeline_type.is_sd_xl() and model_name == "unetxl")
or ((self.version.startswith("flux.1")) and model_name == "transformer")
or (model_name == "unet")
):
config["use_fp8"] = True
elif fp4:
config['use_fp4'] = True
config['onnx_path'] = self._get_onnx_path(model_name, onnx_dir, opt=False, suffix=config['model_suffix'])
config['onnx_opt_path'] = self._get_onnx_path(model_name, onnx_dir, suffix=config['model_suffix'])
config['engine_path'] = self._get_engine_path(model_name, engine_dir, config['do_engine_refit'], suffix=config['model_suffix'])
config['weights_map_path'] = self._get_weights_map_path(model_name, onnx_dir) if config['do_engine_refit'] else None
config['state_dict_path'] = self._get_state_dict_path(model_name, onnx_dir, suffix=config['model_suffix'])
config['refit_weights_path'] = self._get_refit_nodes_path(model_name, onnx_dir, suffix=lora_suffix)
config["onnx_path"] = self.dd_path.model_name_to_unoptimized_onnx_path[model_name]
config["onnx_opt_path"] = self.dd_path.model_name_to_optimized_onnx_path[model_name]
config["engine_path"] = self.dd_path.model_name_to_engine_path[model_name]
config["weights_map_path"] = (
self.dd_path.model_name_to_weights_map_path[model_name] if config["do_engine_refit"] else None
)
config["state_dict_path"] = self.dd_path.model_name_to_quantized_model_state_dict_path[model_name]
config["refit_weights_path"] = self.dd_path.model_name_to_refit_weights_path[model_name]
configs[model_name] = config
@@ -446,7 +445,7 @@ class DiffusionPipeline(ABC):
collate_fn=custom_collate,
)
else:
calibration_file = os.path.join(os.path.dirname(os.path.realpath(__file__)), 'calibration-prompts.txt')
calibration_file = os.path.join(os.getcwd(), 'calibration_data', 'calibration-prompts.txt')
calibration_prompts = load_calib_prompts(calib_batch_size, calibration_file)
def forward_loop(model):
@@ -523,17 +522,44 @@ class DiffusionPipeline(ABC):
return model
def _export_onnx(self, obj, model_config, opt_image_height, opt_image_width, static_shape, onnx_opset, quantization_level, quantization_percentile, quantization_alpha, calibration_size, calib_batch_size, onnx_export_only=False):
@abstractmethod
def download_onnx_models(self, model_name: str, model_config: dict[str, Any]) -> None:
"""Download pre-exported ONNX Models"""
raise NotImplementedError("Please Implement the download_onnx_models method")
def is_native_export_supported(self, model_config: dict[str, Any]) -> bool:
"""Check if pipeline supports native ONNX export"""
# Native export is supported by default
return True
def _export_onnx(
self,
obj,
model_name,
model_config,
opt_image_height,
opt_image_width,
static_shape,
onnx_opset,
quantization_level,
quantization_percentile,
quantization_alpha,
calibration_size,
calib_batch_size,
onnx_export_only,
download_onnx_models,
):
# With onnx_export_only True, the export still happens even if the TRT engine exists. However, it will not re-run the export if the onnx exists.
do_export_onnx = (not os.path.exists(model_config['engine_path']) or onnx_export_only) and not os.path.exists(model_config['onnx_opt_path'])
do_export_weights_map = model_config['weights_map_path'] and not os.path.exists(model_config['weights_map_path'])
if self.version.startswith("flux.1") and model_config['use_fp4']:
# Native export not supported for FP4. Ensure ONNX models exist in the provided directory
assert not do_export_onnx, f"No ONNX model found in {model_config['onnx_opt_path']}. Please download the ONNX models as recommended in the README.md"
if self.version in ["flux.1-dev-canny", "flux.1-dev-depth"] and model_config['use_fp8'] and not self.calibration_dataset:
# Native export of FP8 model requires calibration data. Ensure ONNX models exist in the provided directory
assert not do_export_onnx, f"No ONNX model found in {model_config['onnx_opt_path']}. Please download the ONNX models as recommended in the README.md."
# If ONNX export is required, either download ONNX models or check if the pipeline supports native ONNX export
if do_export_onnx:
if download_onnx_models:
self.download_onnx_models(model_name, model_config)
do_export_onnx = False
else:
self.is_native_export_supported(model_config)
if do_export_onnx or do_export_weights_map:
if not model_config['use_int8'] and not model_config['use_fp8']:
@@ -548,7 +574,6 @@ class DiffusionPipeline(ABC):
print(f"[I] Saving weights map: {model_config['weights_map_path']}")
obj.export_weights_map(model_config['onnx_opt_path'], model_config['weights_map_path'])
def _build_engine(self, obj, engine, model_config, opt_batch_size, opt_image_height, opt_image_width, optimization_level, static_batch, static_shape, enable_all_tactics, timing_cache):
update_output_names = obj.get_output_names() + obj.extra_output_names if obj.extra_output_names else None
fp16amp = False if (model_config['use_fp8'] or getattr(obj, 'build_strongly_typed', False)) else obj.fp16
@@ -556,16 +581,14 @@ class DiffusionPipeline(ABC):
bf16amp = False if (model_config['use_fp8'] or getattr(obj, 'build_strongly_typed', False)) else obj.bf16
strongly_typed = True if (model_config['use_fp8'] or getattr(obj, 'build_strongly_typed', False)) else False
weight_streaming = getattr(obj, 'weight_streaming', False)
extra_build_args = {'verbose': self.verbose}
extra_build_args['builder_optimization_level'] = optimization_level
if model_config['use_int8']:
extra_build_args['int8'] = True
extra_build_args['precision_constraints'] = 'prefer'
int8amp = model_config.get('use_int8', False)
precision_constraints = 'prefer' if int8amp else 'none'
engine.build(model_config['onnx_opt_path'],
strongly_typed=strongly_typed,
fp16=fp16amp,
tf32=tf32amp,
bf16=bf16amp,
int8=int8amp,
input_profile=obj.get_input_profile(
opt_batch_size, opt_image_height, opt_image_width,
static_batch=static_batch, static_shape=static_shape
@@ -575,7 +598,10 @@ class DiffusionPipeline(ABC):
timing_cache=timing_cache,
update_output_names=update_output_names,
weight_streaming=weight_streaming,
**extra_build_args)
verbose=self.verbose,
builder_optimization_level=optimization_level,
precision_constraints=precision_constraints
)
def _refit_engine(self, obj, model_name, model_config):
assert model_config['weights_map_path']
@@ -585,10 +611,12 @@ class DiffusionPipeline(ABC):
if not os.path.exists(model_config['refit_weights_path']):
model = merge_loras(obj.get_model(), self.lora_loader)
refit_weights, updated_weight_names = get_refit_weights(model.state_dict(), model_config['onnx_opt_path'], weights_name_mapping, weights_shape_mapping)
refit_weights, updated_weight_names = engine_module.get_refit_weights(
model.state_dict(), model_config["onnx_opt_path"], weights_name_mapping, weights_shape_mapping
)
print(f"[I] Saving refit weights: {model_config['refit_weights_path']}")
torch.save((refit_weights, updated_weight_names), model_config['refit_weights_path'])
unload_model(model)
torch.save((refit_weights, updated_weight_names), model_config["refit_weights_path"])
unload_torch_model(model)
else:
print(f"[I] Loading refit weights: {model_config['refit_weights_path']}")
refit_weights, updated_weight_names = torch.load(model_config['refit_weights_path'])
@@ -605,9 +633,7 @@ class DiffusionPipeline(ABC):
def load_engines(
self,
engine_dir,
framework_model_dir,
onnx_dir,
onnx_opset,
opt_batch_size,
opt_image_height,
@@ -627,19 +653,15 @@ class DiffusionPipeline(ABC):
calibration_size=32,
calib_batch_size=2,
onnx_export_only=False,
model_onnx_dirs=None,
download_onnx_models=False,
):
"""
Build and load engines for TensorRT accelerated inference.
Export ONNX models first, if applicable.
Args:
engine_dir (str):
Directory to store the TensorRT engines.
framework_model_dir (str):
Directory to store the framework model ckpt.
onnx_dir (str):
Directory to store the ONNX models.
onnx_opset (int):
ONNX opset version to export the models.
opt_batch_size (int):
@@ -679,19 +701,33 @@ class DiffusionPipeline(ABC):
The batch size to use for calibration. Defaults to 2.
onnx_export_only (bool):
Whether only export onnx without building the TRT engine.
model_onnx_dirs (dict(str,str)):
Set onnx dir for each model separately, if not set, use default path in onnx_dir
download_onnx_models (bool):
Download pre-exported ONNX models
"""
self._create_directories(engine_dir, onnx_dir)
self._initialize_models(framework_model_dir, int8, fp8, fp4)
model_configs = self._prepare_model_configs(onnx_dir, model_onnx_dirs, engine_dir, enable_refit, int8, fp8, fp4, quantization_level, quantization_percentile, quantization_alpha, calibration_size)
model_configs = self._prepare_model_configs(enable_refit, int8, fp8, fp4)
# Export models to ONNX
for model_name, obj in self.models.items():
if self.torch_fallback[model_name]:
continue
self._export_onnx(obj, model_configs[model_name], opt_image_height, opt_image_width, static_shape, onnx_opset, quantization_level, quantization_percentile, quantization_alpha, calibration_size, calib_batch_size, onnx_export_only=onnx_export_only)
self._export_onnx(
obj,
model_name,
model_configs[model_name],
opt_image_height,
opt_image_width,
static_shape,
onnx_opset,
quantization_level,
quantization_percentile,
quantization_alpha,
calibration_size,
calib_batch_size,
onnx_export_only,
download_onnx_models,
)
# Release temp GPU memory during onnx export to avoid OOM.
gc.collect()
@@ -706,7 +742,7 @@ class DiffusionPipeline(ABC):
continue
model_config = model_configs[model_name]
engine = Engine(model_config['engine_path'])
engine = engine_module.Engine(model_config["engine_path"])
if not os.path.exists(model_config['engine_path']):
self._build_engine(obj, engine, model_config, opt_batch_size, opt_image_height, opt_image_width, optimization_level, static_batch, static_shape, enable_all_tactics, timing_cache)
self.engine[model_name] = engine
@@ -804,20 +840,18 @@ class DiffusionPipeline(ABC):
prompt_prefix = ''.join(set([prompt[i].replace(' ','_')[:10] for i in range(len(prompt))]))
image_name_prefix = '-'.join([pipeline, prompt_prefix, str(seed)])
image_name_suffix = 'torch' if self.torch_inference else 'trt'
save_image(images, self.output_dir, image_name_prefix, image_name_suffix)
image_module.save_image(images, self.output_dir, image_name_prefix, image_name_suffix)
@abstractmethod
def print_summary(self):
"""Print a summary of the pipeline's configuration."""
raise NotImplementedError("Please Implement the print_summary method")
@abstractmethod
def infer(self):
"""Perform inference using the pipeline."""
raise NotImplementedError("Please Implement the infer method")
@abstractmethod
def run(self):
"""Run the pipeline."""
demo/Diffusion/flux_pipeline.py → demo/Diffusion/demo_diffusion/pipeline/flux_pipeline.py
+162 -38
View File
@@ -1,5 +1,5 @@
#
# SPDX-FileCopyrightText: Copyright (c) 1993-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-FileCopyrightText: Copyright (c) 1993-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#
# Licensed under the Apache License, Version 2.0 (the "License");
@@ -15,29 +15,37 @@
# limitations under the License.
#
import warnings
import numpy as np
from cuda import cudart
import inspect
from models import (
get_clip_embedding_dim,
make_tokenizer,
CLIPModel,
T5Model,
FluxTransformerModel,
VAEModel,
VAEEncoderModel,
)
import tensorrt as trt
import time
import torch
from utilities import (
PIPELINE_TYPE,
TRT_LOGGER,
)
from diffusion_pipeline import DiffusionPipeline
from diffusers.image_processor import VaeImageProcessor
from __future__ import annotations
import argparse
import inspect
import os
import time
import warnings
from typing import Any, List
import numpy as np
import tensorrt as trt
import torch
from cuda import cudart
from diffusers.image_processor import VaeImageProcessor
from huggingface_hub import snapshot_download
from demo_diffusion import path as path_module
from demo_diffusion.model import (
CLIPModel,
FluxTransformerModel,
T5Model,
VAEEncoderModel,
VAEModel,
get_clip_embedding_dim,
load,
make_tokenizer,
)
from demo_diffusion.pipeline.diffusion_pipeline import DiffusionPipeline
from demo_diffusion.pipeline.type import PIPELINE_TYPE
TRT_LOGGER = trt.Logger(trt.Logger.ERROR)
def calculate_shift(
image_seq_len,
@@ -70,7 +78,6 @@ class FluxPipeline(DiffusionPipeline):
weight_streaming=False,
t5_weight_streaming_budget_percentage=None,
transformer_weight_streaming_budget_percentage=None,
force_weakly_typed_t5=False, # NVBug 5071425
**kwargs,
):
"""
@@ -93,18 +100,19 @@ class FluxPipeline(DiffusionPipeline):
transformer_weight_streaming_budget_percentage (`int`, defaults to None):
Weight streaming budget as a percentage of the size of total streamable weights for the transformer model.
"""
super().__init__(version=version, pipeline_type=pipeline_type, weight_streaming=weight_streaming, text_encoder_weight_streaming_budget_percentage=t5_weight_streaming_budget_percentage, denoiser_weight_streaming_budget_percentage=transformer_weight_streaming_budget_percentage, **kwargs)
super().__init__(
version=version,
pipeline_type=pipeline_type,
weight_streaming=weight_streaming,
text_encoder_weight_streaming_budget_percentage=t5_weight_streaming_budget_percentage,
denoiser_weight_streaming_budget_percentage=transformer_weight_streaming_budget_percentage,
**kwargs,
)
self.guidance_scale = guidance_scale
self.max_sequence_length = max_sequence_length
self.bf16=bf16
self.calibration_dataset = calibration_dataset # Currently supported for Flux ControlNet pipelines only
self.low_vram = low_vram
self.force_weakly_typed_t5 = force_weakly_typed_t5
# Pipeline type
self.stages = ["clip", "t5", "transformer", "vae"]
if self.pipeline_type.is_img2img():
self.stages += ["vae_encoder"]
if torch_fallback:
assert type(torch_fallback) is list
@@ -114,6 +122,123 @@ class FluxPipeline(DiffusionPipeline):
self.config[model_name.replace('-','_')+'_torch_fallback'] = True
print(f'[I] Setting torch_fallback for {model_name} model.')
@classmethod
def FromArgs(cls, args: argparse.Namespace, pipeline_type: PIPELINE_TYPE) -> FluxPipeline:
"""Factory method to construct a `FluxPipeline` object from parsed arguments.
Overrides:
DiffusionPipeline.FromArgs
"""
MAX_BATCH_SIZE = 4
DEVICE = "cuda"
DO_RETURN_LATENTS = False
# Resolve all paths.
dd_path = path_module.resolve_path(
cls.get_model_names(pipeline_type), args, pipeline_type, cls._get_pipeline_uid(pipeline_type, args.version)
)
return cls(
dd_path=dd_path,
version=args.version,
pipeline_type=pipeline_type,
guidance_scale=args.guidance_scale,
max_sequence_length=args.max_sequence_length,
bf16=args.bf16,
calibration_dataset=args.calibration_dataset if hasattr(args, "calibration_dataset") else None,
low_vram=args.low_vram,
torch_fallback=args.torch_fallback,
weight_streaming=args.ws,
t5_weight_streaming_budget_percentage=args.t5_ws_percentage,
transformer_weight_streaming_budget_percentage=args.transformer_ws_percentage,
max_batch_size=MAX_BATCH_SIZE,
denoising_steps=args.denoising_steps,
scheduler=args.scheduler,
lora_scale=args.lora_scale,
lora_weight=args.lora_weight,
lora_path=args.lora_path,
device=DEVICE,
output_dir=args.output_dir,
hf_token=args.hf_token,
verbose=args.verbose,
nvtx_profile=args.nvtx_profile,
use_cuda_graph=args.use_cuda_graph,
framework_model_dir=args.framework_model_dir,
return_latents=DO_RETURN_LATENTS,
torch_inference=args.torch_inference,
)
@classmethod
def get_model_names(cls, pipeline_type: PIPELINE_TYPE) -> List[str]:
"""Return a list of model names used by this pipeline.
Overrides:
DiffusionPipeline.get_model_names
"""
if pipeline_type.is_img2img():
return ["clip", "t5", "transformer", "vae", "vae_encoder"]
else:
return ["clip", "t5", "transformer", "vae"]
def download_onnx_models(self, model_name: str, model_config: dict[str, Any]) -> None:
if self.fp16:
raise ValueError(
"ONNX models can be downloaded only for the following precisions: BF16, FP8, FP4. This pipeline is running in FP16."
)
hf_download_path = "-".join([load.get_path(self.version, self.pipeline_type.name), "onnx"])
model_path = model_config["onnx_opt_path"]
base_dir = os.path.dirname(os.path.dirname(model_config["onnx_opt_path"]))
if not os.path.exists(model_path):
if model_name == "clip":
dirname = "clip.opt"
elif model_name == "t5":
dirname = "t5.opt"
elif model_name == "transformer":
if model_config["use_fp4"]:
dirname = "transformer.opt/fp4"
elif model_config["use_fp8"]:
dirname = "transformer.opt/fp8"
elif self.bf16:
dirname = "transformer.opt/bf16"
elif model_name == "vae":
dirname = "vae.opt"
elif model_name == "vae_encoder":
dirname = "vae_encoder.opt"
else:
raise ValueError(f"{model_name} not found in {self.stages}")
snapshot_download(
repo_id=hf_download_path,
allow_patterns=os.path.join(dirname, "*"),
local_dir=base_dir,
token=self.hf_token,
)
# Rename directory from <model_name>.opt to <model_name>
saved_dir = os.path.join(base_dir, dirname)
model_dir = os.path.dirname(model_path)
os.rename(saved_dir, model_dir)
# Rename model from model.onnx to model_optimized.onnx
os.rename(os.path.join(model_dir, "model.onnx"), model_path)
def is_native_export_supported(self, model_config: dict[str, Any]) -> bool:
if self.version.startswith("flux.1") and model_config["use_fp4"]:
# Native export not supported for FP4.
raise ValueError(
f"No ONNX model found in {model_config['onnx_opt_path']}. Please pass --download-onnx-models."
)
if (
self.version in ["flux.1-dev-canny", "flux.1-dev-depth"]
and model_config["use_fp8"]
and not self.calibration_dataset
):
# Native export of FP8 model requires calibration data.
raise ValueError(
f"No ONNX model found in {model_config['onnx_opt_path']}. Please pass --download-onnx-models. If you would like to quantize and export natively, please provide calibration data using --calibration-."
)
return True
def _initialize_models(self, framework_model_dir, int8, fp8, fp4):
# Load text tokenizer(s)
self.tokenizer = make_tokenizer(
@@ -162,7 +287,7 @@ class FluxPipeline(DiffusionPipeline):
bf16=self.bf16,
subfolder="text_encoder_2",
text_maxlen=self.max_sequence_length,
build_strongly_typed=False if self.force_weakly_typed_t5 and self.fp16 else True,
build_strongly_typed=True,
weight_streaming=self.weight_streaming,
weight_streaming_budget_percentage=self.text_encoder_weight_streaming_budget_percentage,
)
@@ -170,8 +295,8 @@ class FluxPipeline(DiffusionPipeline):
if "transformer" in self.stages:
self.models["transformer"] = FluxTransformerModel(
**models_args,
bf16=True if int8 or fp8 else self.bf16,
fp16=False if int8 or fp8 else self.fp16,
bf16=self.bf16,
fp16=self.fp16,
int8=int8,
fp8=fp8,
tf32=self.tf32,
@@ -183,7 +308,7 @@ class FluxPipeline(DiffusionPipeline):
if "vae" in self.stages:
# Accuracy issues with FP16
self.models["vae"] = VAEModel(**models_args, fp16=False, tf32=self.tf32, bf16=True if int8 or fp8 else self.bf16)
self.models["vae"] = VAEModel(**models_args, fp16=False, tf32=self.tf32, bf16=self.bf16)
self.vae_scale_factor = (
2 ** (len(self.models["vae"].config["block_out_channels"]))
@@ -191,9 +316,8 @@ class FluxPipeline(DiffusionPipeline):
else 16
)
if 'vae_encoder' in self.stages:
# WAR: VAE Encoder fallback to FP32 in BF16 TRT pipeline. TRT support will be added in a future release
self.models['vae_encoder'] = VAEEncoderModel(**models_args, fp16=False, tf32=self.tf32, bf16=self.bf16 if self.torch_inference else False)
if "vae_encoder" in self.stages:
self.models['vae_encoder'] = VAEEncoderModel(**models_args, fp16=False, tf32=self.tf32, bf16=self.bf16)
self.vae_latent_channels = (
self.models["vae"].config["latent_channels"] if "vae" in self.stages and self.models["vae"] is not None else 16
)
demo/Diffusion/stable_cascade_pipeline.py → demo/Diffusion/demo_diffusion/pipeline/stable_cascade_pipeline.py
+17 -17
View File
@@ -1,5 +1,5 @@
#
# SPDX-FileCopyrightText: Copyright (c) 1993-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-FileCopyrightText: Copyright (c) 1993-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#
# Licensed under the Apache License, Version 2.0 (the "License");
@@ -15,25 +15,25 @@
# limitations under the License.
#
from cuda import cudart
from diffusers import (
DDPMWuerstchenScheduler
)
import inspect
from models import (
make_tokenizer,
import time
import tensorrt as trt
import torch
from cuda import cudart
from diffusers import DDPMWuerstchenScheduler
from demo_diffusion.model import (
CLIPWithProjModel,
UNetCascadeModel,
VQGANModel
VQGANModel,
make_tokenizer,
)
import tensorrt as trt
import time
import torch
from utilities import (
PIPELINE_TYPE,
TRT_LOGGER,
)
from stable_diffusion_pipeline import StableDiffusionPipeline
from demo_diffusion.pipeline.stable_diffusion_pipeline import StableDiffusionPipeline
from demo_diffusion.pipeline.type import PIPELINE_TYPE
TRT_LOGGER = trt.Logger(trt.Logger.ERROR)
class StableCascadePipeline(StableDiffusionPipeline):
"""
@@ -281,7 +281,7 @@ class StableCascadePipeline(StableDiffusionPipeline):
# TODO: support custom timesteps
timesteps = None
if timesteps is not None:
if not ("timesteps" in set(inspect.signature(self.scheduler.set_timesteps).parameters.keys())):
if "timesteps" not in set(inspect.signature(self.scheduler.set_timesteps).parameters.keys()):
raise ValueError(
f"The current scheduler class {self.scheduler.__class__}'s `set_timesteps` does not support custom"
f" timestep schedules. Please check whether you are using the correct scheduler."
demo/Diffusion/stable_diffusion_3_pipeline.py → demo/Diffusion/demo_diffusion/pipeline/stable_diffusion_3_pipeline.py
+31 -28
View File
@@ -1,5 +1,5 @@
#
# SPDX-FileCopyrightText: Copyright (c) 1993-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-FileCopyrightText: Copyright (c) 1993-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#
# Licensed under the Apache License, Version 2.0 (the "License");
@@ -14,32 +14,33 @@
# See the License for the specific language governing permissions and
# limitations under the License.
#
import math
import os
import pathlib
import time
import nvtx
import tensorrt as trt
import torch
from cuda import cudart
from models import (
get_clip_embedding_dim,
import demo_diffusion.engine as engine_module
import demo_diffusion.image as image_module
from demo_diffusion.model import (
SD3_CLIPGModel,
SD3_CLIPLModel,
SD3_T5XXLModel,
SD3_MMDiTModel,
SD3_T5XXLModel,
SD3_VAEDecoderModel,
SD3_VAEEncoderModel,
SD3_VAEDecoderModel
get_clip_embedding_dim,
)
import nvtx
import os
import math
import pathlib
import tensorrt as trt
import time
import torch
from utilities import (
PIPELINE_TYPE,
TRT_LOGGER,
Engine,
save_image,
)
from utils_sd3.other_impls import SD3Tokenizer
from utils_sd3.sd3_impls import SD3LatentFormat, sample_euler
from demo_diffusion.pipeline.type import PIPELINE_TYPE
from demo_diffusion.utils_sd3.other_impls import SD3Tokenizer
from demo_diffusion.utils_sd3.sd3_impls import SD3LatentFormat, sample_euler
TRT_LOGGER = trt.Logger(trt.Logger.ERROR)
class StableDiffusion3Pipeline:
"""
@@ -206,6 +207,7 @@ class StableDiffusion3Pipeline:
static_shape=True,
enable_all_tactics=False,
timing_cache=None,
**_kwargs,
):
"""
Build and load engines for TensorRT accelerated inference.
@@ -292,7 +294,7 @@ class StableDiffusion3Pipeline:
for model_name, obj in self.models.items():
if self.torch_fallback[model_name]:
continue
engine = Engine(engine_path[model_name])
engine = engine_module.Engine(engine_path[model_name])
if not os.path.exists(engine_path[model_name]):
update_output_names = obj.get_output_names() + obj.extra_output_names if obj.extra_output_names else None
extra_build_args = {'verbose': self.verbose}
@@ -306,7 +308,8 @@ class StableDiffusion3Pipeline:
enable_all_tactics=enable_all_tactics,
timing_cache=timing_cache,
update_output_names=update_output_names,
**extra_build_args)
verbose=self.verbose
)
self.engine[model_name] = engine
# Load TensorRT engines
@@ -374,7 +377,7 @@ class StableDiffusion3Pipeline:
# Save image
image_name_prefix = pipeline+''.join(set(['-'+prompt[i].replace(' ','_')[:10] for i in range(len(prompt))]))+'-'+str(seed)+'-'
image_name_suffix = 'torch' if self.torch_inference else 'trt'
save_image(images, self.output_dir, image_name_prefix, image_name_suffix)
image_module.save_image(images, self.output_dir, image_name_prefix, image_name_suffix)
def encode_prompt(self, prompt, negative_prompt):
def encode_token_weights(model_name, token_weight_pairs):
@@ -387,7 +390,7 @@ class StableDiffusion3Pipeline:
else:
trt_out = self.runEngine(model_name, {'input_ids': tokens})
out, pooled = trt_out['text_embeddings'], trt_out["pooled_output"]
self.profile_stop(model_name)
if pooled is not None:
@@ -410,7 +413,7 @@ class StableDiffusion3Pipeline:
conditioning = tokenize(prompt[0])
neg_conditioning = tokenize(negative_prompt[0])
return conditioning, neg_conditioning
def denoise_latent(self, latent, conditioning, neg_conditioning, model_name='mmdit'):
def get_noise(latent):
return torch.randn(latent.size(), dtype=torch.float32, layout=latent.layout, generator=self.generator, device="cuda").to(latent.dtype)
@@ -434,7 +437,7 @@ class StableDiffusion3Pipeline:
def fix_cond(cond):
cond, pooled = (cond[0].half().cuda(), cond[1].half().cuda())
return { "c_crossattn": cond, "y": pooled }
def cfg_denoiser(x, timestep, cond, uncond, cond_scale):
# Run cond and uncond in a batch together
sample = torch.cat([x, x])
@@ -540,7 +543,7 @@ class StableDiffusion3Pipeline:
unet_channels=16,
latent_height=latent_height,
latent_width=latent_width)
# Encode input image
if self.input_image is not None:
latent = self.encode_image()
@@ -550,7 +553,7 @@ class StableDiffusion3Pipeline:
# Denoise
latent = self.denoise_latent(latent, conditioning, neg_conditioning)
# Decode Latents
images = self.decode_latent(latent)
demo/Diffusion/stable_diffusion_pipeline.py → demo/Diffusion/demo_diffusion/pipeline/stable_diffusion_pipeline.py
+80 -56
View File
@@ -1,5 +1,5 @@
#
# SPDX-FileCopyrightText: Copyright (c) 1993-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-FileCopyrightText: Copyright (c) 1993-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#
# Licensed under the Apache License, Version 2.0 (the "License");
@@ -15,64 +15,66 @@
# limitations under the License.
#
import inspect
import json
import os
import pathlib
import time
from hashlib import md5
from typing import List, Optional
import modelopt.torch.opt as mto
import modelopt.torch.quantization as mtq
import numpy as np
import nvtx
import tensorrt as trt
import torch
from cuda import cudart
from diffusers import (
DDIMScheduler,
DDPMScheduler,
EulerDiscreteScheduler,
DDPMWuerstchenScheduler,
EulerAncestralDiscreteScheduler,
LCMScheduler, LMSDiscreteScheduler,
EulerDiscreteScheduler,
LCMScheduler,
LMSDiscreteScheduler,
PNDMScheduler,
UniPCMultistepScheduler,
DDPMWuerstchenScheduler
)
from hashlib import md5
import inspect
from models import (
get_clip_embedding_dim,
LoraLoader,
make_tokenizer,
import demo_diffusion.engine as engine_module
import demo_diffusion.image as image_module
from demo_diffusion.model import (
CLIPModel,
CLIPWithProjModel,
LoraLoader,
UNetModel,
UNetXLModel,
VAEModel,
UNetXLModelControlNet,
VAEEncoderModel,
make_scheduler
)
import numpy as np
import nvtx
import json
import os
import pathlib
import tensorrt as trt
import time
import torch
from typing import Optional, List
from utilities import (
PIPELINE_TYPE,
TRT_LOGGER,
Engine,
get_refit_weights,
load_calib_prompts,
VAEModel,
get_clip_embedding_dim,
make_scheduler,
make_tokenizer,
merge_loras,
prepare_mask_and_masked_image,
save_image,
unload_model
unload_torch_model,
)
from utils_modelopt import (
filter_func,
quantize_lvl,
get_int8_config,
check_lora,
set_fmha,
generate_fp8_scales,
from demo_diffusion.pipeline.calibrate import load_calib_prompts
from demo_diffusion.pipeline.type import PIPELINE_TYPE
from demo_diffusion.utils_modelopt import (
SD_FP8_FP16_DEFAULT_CONFIG,
SD_FP8_FP32_DEFAULT_CONFIG,
check_lora,
filter_func,
generate_fp8_scales,
get_int8_config,
quantize_lvl,
set_fmha,
)
TRT_LOGGER = trt.Logger(trt.Logger.ERROR)
class StableDiffusionPipeline:
SCHEDULER_DEFAULTS = {
"1.4": "PNDM",
@@ -179,7 +181,7 @@ class StableDiffusionPipeline:
# Pipeline type
self.pipeline_type = pipeline_type
if self.pipeline_type.is_txt2img() or self.pipeline_type.is_controlnet():
if self.pipeline_type.is_txt2img():
self.stages = ['clip','unet','vae']
elif self.pipeline_type.is_img2img() or self.pipeline_type.is_inpaint():
self.stages = ['vae_encoder', 'clip','unet','vae']
@@ -339,7 +341,23 @@ class StableDiffusionPipeline:
self.models['unet'] = UNetModel(**models_args, fp16=True, int8=int8, fp8=fp8, controlnets=self.controlnets, do_classifier_free_guidance=self.do_classifier_free_guidance)
if 'unetxl' in self.stages:
self.models['unetxl'] = UNetXLModel(**models_args, fp16=True, int8=int8, fp8=fp8, do_classifier_free_guidance=self.do_classifier_free_guidance)
if not self.controlnets:
self.models["unetxl"] = UNetXLModel(
**models_args,
fp16=True,
int8=int8,
fp8=fp8,
do_classifier_free_guidance=self.do_classifier_free_guidance,
)
else:
self.models["unetxl"] = UNetXLModelControlNet(
**models_args,
fp16=True,
int8=int8,
fp8=fp8,
controlnets=self.controlnets,
do_classifier_free_guidance=self.do_classifier_free_guidance,
)
vae_fp16 = not self.pipeline_type.is_sd_xl()
@@ -371,6 +389,7 @@ class StableDiffusionPipeline:
quantization_alpha=0.8,
calibration_size=32,
calib_batch_size=2,
**_kwargs,
):
"""
Build and load engines for TensorRT accelerated inference.
@@ -475,7 +494,7 @@ class StableDiffusionPipeline:
state_dict_path = self.getStateDictPath(model_name, onnx_dir, suffix=model_suffix[model_name])
if not os.path.exists(state_dict_path):
print(f"[I] Calibrated weights not found, generating {state_dict_path}")
calibration_file = os.path.join(os.path.dirname(os.path.realpath(__file__)), 'calibration-prompts.txt')
calibration_file = os.path.join(os.getcwd(), 'calibration_data', 'calibration-prompts.txt')
calibration_prompts = load_calib_prompts(calib_batch_size, calibration_file)
# TODO check size > calibration_size
def do_calibrate(pipeline, calibration_prompts, **kwargs):
@@ -536,7 +555,7 @@ class StableDiffusionPipeline:
for model_name, obj in self.models.items():
if torch_fallback[model_name]:
continue
engine = Engine(engine_path[model_name])
engine = engine_module.Engine(engine_path[model_name])
if not os.path.exists(engine_path[model_name]):
update_output_names = obj.get_output_names() + obj.extra_output_names if obj.extra_output_names else None
fp16amp = obj.fp16 if not use_fp8[model_name] else False
@@ -544,16 +563,14 @@ class StableDiffusionPipeline:
# TF32 can be enabled for all precisions (including INT8/FP8)
tf32amp = obj.tf32
strongly_typed = False if not use_fp8[model_name] else True
extra_build_args = {'verbose': self.verbose}
extra_build_args['builder_optimization_level'] = optimization_level
if use_int8[model_name]:
extra_build_args['int8'] = True
extra_build_args['precision_constraints'] = 'prefer'
int8amp = use_int8.get('model_name', False)
precision_constraints = 'prefer' if int8amp else 'none'
engine.build(onnx_opt_path[model_name],
strongly_typed=strongly_typed,
fp16=fp16amp,
bf16=bf16amp,
tf32=tf32amp,
int8=int8amp,
input_profile=obj.get_input_profile(
opt_batch_size, opt_image_height, opt_image_width,
static_batch=static_batch, static_shape=static_shape
@@ -562,7 +579,10 @@ class StableDiffusionPipeline:
enable_all_tactics=enable_all_tactics,
timing_cache=timing_cache,
update_output_names=update_output_names,
**extra_build_args)
verbose=self.verbose,
builder_optimization_level=optimization_level,
precision_constraints=precision_constraints,
)
self.engine[model_name] = engine
# Load TensorRT engines
@@ -577,11 +597,13 @@ class StableDiffusionPipeline:
[weights_name_mapping, weights_shape_mapping] = json.load(fp_wts)
refit_weights_path = self.getRefitNodesPath(model_name, engine_dir, suffix=lora_suffix)
if not os.path.exists(refit_weights_path):
print(f"[I] Saving refit weights: {refit_weights_path}")
model = merge_loras(obj.get_model(), self.lora_loader)
refit_weights, updated_weight_names = get_refit_weights(model.state_dict(), onnx_opt_path[model_name], weights_name_mapping, weights_shape_mapping)
torch.save((refit_weights, updated_weight_names), refit_weights_path)
unload_model(model)
print(f"[I] Saving refit weights: {refit_weights_path}")
model = merge_loras(obj.get_model(), self.lora_loader)
refit_weights, updated_weight_names = engine_module.get_refit_weights(
model.state_dict(), onnx_opt_path[model_name], weights_name_mapping, weights_shape_mapping
)
torch.save((refit_weights, updated_weight_names), refit_weights_path)
unload_torch_model(model)
else:
print(f"[I] Loading refit weights: {refit_weights_path}")
refit_weights, updated_weight_names = torch.load(refit_weights_path)
@@ -861,7 +883,7 @@ class StableDiffusionPipeline:
# Save image
image_name_prefix = pipeline+''.join(set(['-'+prompt[i].replace(' ','_')[:10] for i in range(len(prompt))]))+'-'+str(seed)+'-'
image_name_suffix = 'torch' if self.torch_inference else 'trt'
save_image(images, self.output_dir, image_name_prefix, image_name_suffix)
image_module.save_image(images, self.output_dir, image_name_prefix, image_name_suffix)
def infer(
self,
@@ -926,7 +948,7 @@ class StableDiffusionPipeline:
# TODO: support custom timesteps
timesteps = None
if timesteps is not None:
if not ("timesteps" in set(inspect.signature(self.scheduler.set_timesteps).parameters.keys())):
if "timesteps" not in set(inspect.signature(self.scheduler.set_timesteps).parameters.keys()):
raise ValueError(
f"The current scheduler class {self.scheduler.__class__}'s `set_timesteps` does not support custom"
f" timestep schedules. Please check whether you are using the correct scheduler."
@@ -962,7 +984,9 @@ class StableDiffusionPipeline:
noise = torch.randn(image_latents.shape, generator=self.generator, device=self.device, dtype=torch.float32)
latents = self.scheduler.add_noise(image_latents, noise, latent_timestep)
elif self.pipeline_type.is_inpaint():
mask, mask_image = self.preprocess_images(batch_size, prepare_mask_and_masked_image(input_image, mask_image))
mask, mask_image = self.preprocess_images(
batch_size, image_module.prepare_mask_and_masked_image(input_image, mask_image)
)
mask = torch.nn.functional.interpolate(mask, size=(latent_height, latent_width))
mask = torch.cat([mask] * 2)
masked_image_latents = self.encode_image(mask_image)
demo/Diffusion/stable_video_diffusion_pipeline.py → demo/Diffusion/demo_diffusion/pipeline/stable_video_diffusion_pipeline.py
+48 -39
View File
@@ -1,6 +1,6 @@
#
# Copyright 2024 The HuggingFace Inc. team.
# SPDX-FileCopyrightText: Copyright (c) 1993-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-FileCopyrightText: Copyright (c) 1993-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#
# Licensed under the Apache License, Version 2.0 (the "License");
@@ -16,45 +16,54 @@
# limitations under the License.
#
import os
import pathlib
import random
import time
from typing import Optional
import modelopt.torch.opt as mto
import modelopt.torch.quantization as mtq
import tensorrt as trt
import torch
from cuda import cudart
from diffusers.image_processor import VaeImageProcessor
from diffusers.utils.torch_utils import randn_tensor
from models import (
from tqdm.auto import tqdm
import demo_diffusion.engine as engine_module
import demo_diffusion.image as image_module
from demo_diffusion.model import (
CLIPImageProcessorModel,
CLIPVisionWithProjModel,
UNetTemporalModel,
VAEDecTemporalModel,
)
import os
import pathlib
import tensorrt as trt
import time
import torch
import random
from tqdm.auto import tqdm
from typing import Optional
from utilities import (
GiB,
PIPELINE_TYPE,
TRT_LOGGER,
Engine,
_append_dims,
_resize_with_antialiasing,
tensor2vid,
load_calibration_images,
)
import modelopt.torch.opt as mto
import modelopt.torch.quantization as mtq
from utils_modelopt import (
filter_func,
quantize_lvl,
check_lora,
set_fmha,
generate_fp8_scales,
from demo_diffusion.pipeline.calibrate import load_calibration_images
from demo_diffusion.pipeline.stable_diffusion_pipeline import StableDiffusionPipeline
from demo_diffusion.pipeline.type import PIPELINE_TYPE
from demo_diffusion.utils_modelopt import (
SD_FP8_FP16_DEFAULT_CONFIG,
check_lora,
filter_func,
generate_fp8_scales,
quantize_lvl,
)
from stable_diffusion_pipeline import StableDiffusionPipeline
TRT_LOGGER = trt.Logger(trt.Logger.ERROR)
def _GiB(val):
return val * 1 << 30
def _append_dims(x, target_dims):
"""Appends dimensions to the end of a tensor until it has target_dims dimensions."""
dims_to_append = target_dims - x.ndim
if dims_to_append < 0:
raise ValueError(f"input has {x.ndim} dims but target_dims is {target_dims}, which is less")
return x[(...,) + (None,) * dims_to_append]
class StableVideoDiffusionPipeline(StableDiffusionPipeline):
"""
@@ -110,11 +119,11 @@ class StableVideoDiffusionPipeline(StableDiffusionPipeline):
# TODO user configurable cuda_device_id
cuda_device_id = 0
vram_size = cudart.cudaGetDeviceProperties(cuda_device_id)[1].totalGlobalMem
self.low_vram = vram_size < GiB(40)
self.low_vram = vram_size < _GiB(40)
if self.low_vram:
print(f"[W] WARNING low VRAM ({vram_size/GiB(1):.2f} GB) mode selected. Certain optimizations may be skipped.")
print(f"[W] WARNING low VRAM ({vram_size/_GiB(1):.2f} GB) mode selected. Certain optimizations may be skipped.")
if self.use_cuda_graph and self.low_vram:
print(f"[W] WARNING CUDA graph disabled in low VRAM mode.")
print("[W] WARNING CUDA graph disabled in low VRAM mode.")
self.use_cuda_graph = False
self.config = {}
@@ -166,7 +175,8 @@ class StableVideoDiffusionPipeline(StableDiffusionPipeline):
fp8=False,
quantization_level=0.0,
calibration_size=32,
calib_batch_size=2
calib_batch_size=2,
**_kwargs,
):
"""
Build and load engines for TensorRT accelerated inference.
@@ -246,7 +256,6 @@ class StableVideoDiffusionPipeline(StableDiffusionPipeline):
engine_path = { model_name : self.getEnginePath(model_name, engine_dir, do_engine_refit[model_name], suffix=model_suffix[model_name]) for model_name in model_names }
weights_map_path = { model_name : (self.getWeightsMapPath(model_name, onnx_dir) if do_engine_refit[model_name] else None) for model_name in model_names }
# Export models to ONNX
for model_name, obj in self.models.items():
if self.torch_fallback[model_name]:
@@ -262,7 +271,7 @@ class StableVideoDiffusionPipeline(StableDiffusionPipeline):
if not os.path.exists(state_dict_path):
# Load calibration images
print(f"[I] Calibrated weights not found, generating {state_dict_path}")
calibration_image_folder = os.path.join(os.path.dirname(os.path.realpath(__file__)), 'calibration-images')
calibration_image_folder = os.path.join(os.getcwd(), 'calibration_data', 'calibration-images')
calibration_image_list = load_calibration_images(calibration_image_folder)
print("Number of images loaded:", len(calibration_image_list))
@@ -319,7 +328,7 @@ class StableVideoDiffusionPipeline(StableDiffusionPipeline):
for model_name, obj in self.models.items():
if self.torch_fallback[model_name]:
continue
engine = Engine(engine_path[model_name])
engine = engine_module.Engine(engine_path[model_name])
if not os.path.exists(engine_path[model_name]):
update_output_names = obj.get_output_names() + obj.extra_output_names if obj.extra_output_names else None
engine.build(onnx_opt_path[model_name],
@@ -359,7 +368,7 @@ class StableVideoDiffusionPipeline(StableDiffusionPipeline):
self.engine[model_name].allocate_buffers(shape_dict=alloc_shape, device=self.device)
else:
# Load torch model
if not model_name in self.torch_models:
if model_name not in self.torch_models:
self.torch_models[model_name] = self.models[model_name].get_model(torch_inference=self.torch_inference)
def deactivateEngines(self, model_name, release_model=True):
@@ -406,7 +415,7 @@ class StableVideoDiffusionPipeline(StableDiffusionPipeline):
# We normalize the image before resizing to match with the original implementation.
# Then we unnormalize it after resizing.
image = image * 2.0 - 1.0
image = _resize_with_antialiasing(image, (224, 224))
image = image_module.resize_with_antialiasing(image, (224, 224))
image = (image + 1.0) / 2.0
# Normalize the image with for CLIP input
@@ -706,7 +715,7 @@ class StableVideoDiffusionPipeline(StableDiffusionPipeline):
self.profile_start('vae', color='red')
self.torch_models['vae-temp'].to(dtype=torch.float16)
frames = self.decode_latents(latents, self.num_frames, self.decode_chunk_size)
frames = tensor2vid(frames, self.image_processor, output_type='pil')
frames = image_module.tensor2vid(frames, self.image_processor, output_type="pil")
self.profile_stop('vae')
torch.cuda.synchronize()
demo/Diffusion/demo_diffusion/pipeline/type.py
+64
View File
@@ -0,0 +1,64 @@
#
# SPDX-FileCopyrightText: Copyright (c) 1993-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
import enum
class PIPELINE_TYPE(enum.Enum):
TXT2IMG = enum.auto()
IMG2IMG = enum.auto()
IMG2VID = enum.auto()
INPAINT = enum.auto()
CONTROLNET = enum.auto()
XL_CONTROLNET = enum.auto()
XL_BASE = enum.auto()
XL_REFINER = enum.auto()
CASCADE_PRIOR = enum.auto()
CASCADE_DECODER = enum.auto()
def is_txt2img(self):
return self in (self.TXT2IMG, self.CONTROLNET)
def is_img2img(self):
return self == self.IMG2IMG
def is_img2vid(self):
return self == self.IMG2VID
def is_inpaint(self):
return self == self.INPAINT
def is_controlnet(self):
return self in (self.CONTROLNET, self.XL_CONTROLNET)
def is_sd_xl_base(self):
return self in (self.XL_BASE, self.XL_CONTROLNET)
def is_sd_xl_refiner(self):
return self == self.XL_REFINER
def is_sd_xl(self):
return self.is_sd_xl_base() or self.is_sd_xl_refiner()
def is_cascade_prior(self):
return self == self.CASCADE_PRIOR
def is_cascade_decoder(self):
return self == self.CASCADE_DECODER
def is_cascade(self):
return self.is_cascade_prior() or self.is_cascade_decoder()
demo/Diffusion/utils_modelopt.py → demo/Diffusion/demo_diffusion/utils_modelopt.py
Regular → Executable
+1 -1
View File
@@ -1,5 +1,5 @@
#
# SPDX-FileCopyrightText: Copyright (c) 1993-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-FileCopyrightText: Copyright (c) 1993-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#
# Licensed under the Apache License, Version 2.0 (the "License");
demo/Diffusion/demo_diffusion/utils_sd3/__init__.py
View File
demo/Diffusion/utils_sd3/mmdit.py → demo/Diffusion/demo_diffusion/utils_sd3/mmdit.py
+4 -1
View File
@@ -22,11 +22,14 @@
import math
from typing import Dict, Optional
import numpy as np
import torch
import torch.nn as nn
from einops import rearrange, repeat
from utils_sd3.other_impls import attention, Mlp
from demo_diffusion.utils_sd3.other_impls import Mlp, attention
class PatchEmbed(nn.Module):
""" 2D Image to Patch Embedding"""
demo/Diffusion/utils_sd3/other_impls.py → demo/Diffusion/demo_diffusion/utils_sd3/other_impls.py
View File
demo/Diffusion/utils_sd3/sd3_impls.py → demo/Diffusion/demo_diffusion/utils_sd3/sd3_impls.py
+5 -2
View File
@@ -20,10 +20,13 @@
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.
import torch, math, einops
from utils_sd3.mmdit import MMDiT
import math
import einops
import torch
from PIL import Image
from demo_diffusion.utils_sd3.mmdit import MMDiT
#################################################################################################
### MMDiT Model Wrapping
demo/Diffusion/demo_img2img.py
+12 -17
View File
@@ -1,5 +1,5 @@
#
# SPDX-FileCopyrightText: Copyright (c) 1993-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-FileCopyrightText: Copyright (c) 1993-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#
# Licensed under the Apache License, Version 2.0 (the "License");
@@ -21,19 +21,14 @@ import PIL
from cuda import cudart
from PIL import Image
from stable_diffusion_pipeline import StableDiffusionPipeline
from utilities import (
PIPELINE_TYPE,
TRT_LOGGER,
add_arguments,
download_image,
preprocess_image,
process_pipeline_args
)
from demo_diffusion import dd_argparse
from demo_diffusion import image as image_module
from demo_diffusion import pipeline as pipeline_module
def parseArgs():
parser = argparse.ArgumentParser(description="Options for Stable Diffusion Img2Img Demo")
parser = add_arguments(parser)
parser = dd_argparse.add_arguments(parser)
parser.add_argument('--input-image', type=str, default="", help="Path to the input image")
return parser.parse_args()
@@ -45,7 +40,7 @@ if __name__ == "__main__":
input_image = Image.open(args.input_image)
else:
url = "https://raw.githubusercontent.com/CompVis/stable-diffusion/main/assets/stable-samples/img2img/sketch-mountains-input.jpg"
input_image = download_image(url)
input_image = image_module.download_image(url)
image_width, image_height = input_image.size
if image_height != args.height or image_width != args.width:
@@ -54,14 +49,14 @@ if __name__ == "__main__":
image_height, image_width = args.height, args.width
if isinstance(input_image, PIL.Image.Image):
input_image = preprocess_image(input_image)
input_image = image_module.preprocess_image(input_image)
kwargs_init_pipeline, kwargs_load_engine, args_run_demo = process_pipeline_args(args)
kwargs_init_pipeline, kwargs_load_engine, args_run_demo = dd_argparse.process_pipeline_args(args)
# Initialize demo
demo = StableDiffusionPipeline(
pipeline_type=PIPELINE_TYPE.IMG2IMG,
**kwargs_init_pipeline)
demo = pipeline_module.StableDiffusionPipeline(
pipeline_type=pipeline_module.PIPELINE_TYPE.IMG2IMG, **kwargs_init_pipeline
)
# Load TensorRT engines and pytorch modules
demo.loadEngines(
demo/Diffusion/demo_img2img_flux.py
+13 -51
View File
@@ -1,5 +1,5 @@
#
# SPDX-FileCopyrightText: Copyright (c) 1993-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-FileCopyrightText: Copyright (c) 1993-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#
# Licensed under the Apache License, Version 2.0 (the "License");
@@ -20,21 +20,15 @@ import os
import controlnet_aux
from cuda import cudart
from flux_pipeline import FluxPipeline
from PIL import Image
from utilities import (
PIPELINE_TYPE,
VALID_OPTIMIZATION_LEVELS,
add_arguments,
process_pipeline_args,
)
from demo_diffusion import dd_argparse
from demo_diffusion import pipeline as pipeline_module
def parse_args():
parser = argparse.ArgumentParser(
description="Options for Flux Img2Img Demo", conflict_handler="resolve"
)
parser = add_arguments(parser)
parser = argparse.ArgumentParser(description="Options for Flux Img2Img Demo", conflict_handler="resolve")
parser = dd_argparse.add_arguments(parser)
parser.add_argument(
"--version",
type=str,
@@ -82,7 +76,7 @@ def parse_args():
"--optimization-level",
type=int,
default=3,
help=f"Set the builder optimization level to build the engine with. A higher level allows TensorRT to spend more building time for more optimization options. Must be one of {VALID_OPTIMIZATION_LEVELS}.",
help=f"Set the builder optimization level to build the engine with. A higher level allows TensorRT to spend more building time for more optimization options. Must be one of {dd_argparse.VALID_OPTIMIZATION_LEVELS}.",
)
parser.add_argument(
"--torch-fallback",
@@ -120,25 +114,6 @@ def parse_args():
help="If set, only performs the export of models to ONNX, skipping engine build and inference.",
)
def _parse_kv_pairs(s: str):
"""Parse a string of key-value pairs into a dictionary.
Expected format: key1:value1,key2:value2,...
"""
result = {}
# Split by comma to get each pair
pairs = s.split(",")
for pair in pairs:
# Split by ':' to separate key and value
key, value = pair.split(":", 1)
result[key] = value
return result
parser.add_argument(
"--model-onnx-dirs",
type=_parse_kv_pairs,
help="Set directories for individual ONNX models. For example: --model-onnx-dirs=transformer:/path/to/transformer,vae:/path/to/vae,t5:/path/to/t5,clip:/path/to/clip",
)
parser.add_argument(
"--calibration-dataset",
type=str,
@@ -183,7 +158,9 @@ def process_demo_args(args):
args.torch_fallback = args.torch_fallback.split(",")
if args.torch_fallback and args.torch_inference:
print("[W] All models will run in PyTorch when --torch-inference is set. Parameter --torch-fallback will be ignored.")
print(
"[W] All models will run in PyTorch when --torch-inference is set. Parameter --torch-fallback will be ignored."
)
args.torch_fallback = None
controlnet_type = "depth" if "depth" in args.version else "canny" if "canny" in args.version else ""
@@ -253,31 +230,16 @@ if __name__ == "__main__":
print("[I] Initializing Flux img2img demo using TensorRT")
args = parse_args()
kwargs_init_pipeline, kwargs_load_engine, _ = process_pipeline_args(args)
_, kwargs_load_engine, _ = dd_argparse.process_pipeline_args(args)
args_run_demo = process_demo_args(args)
# Initialize demo
demo = FluxPipeline(
pipeline_type=PIPELINE_TYPE.IMG2IMG,
max_sequence_length=args.max_sequence_length,
bf16=args.bf16,
calibration_dataset=args.calibration_dataset,
low_vram=args.low_vram,
torch_fallback=args.torch_fallback,
weight_streaming=args.ws,
t5_weight_streaming_budget_percentage=args.t5_ws_percentage,
transformer_weight_streaming_budget_percentage=args.transformer_ws_percentage,
**kwargs_init_pipeline,
)
demo = pipeline_module.FluxPipeline.FromArgs(args, pipeline_type=pipeline_module.PIPELINE_TYPE.IMG2IMG)
# Load TensorRT engines and pytorch modules
demo.load_engines(
args.engine_dir,
args.framework_model_dir,
args.onnx_dir,
framework_model_dir=args.framework_model_dir,
onnx_export_only=args.onnx_export_only,
model_onnx_dirs=args.model_onnx_dirs,
fp4=args.fp4,
**kwargs_load_engine,
)
demo/Diffusion/demo_img2vid.py
+11 -13
View File
@@ -1,5 +1,5 @@
#
# SPDX-FileCopyrightText: Copyright (c) 1993-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-FileCopyrightText: Copyright (c) 1993-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#
# Licensed under the Apache License, Version 2.0 (the "License");
@@ -19,16 +19,14 @@ import argparse
from PIL import Image
from stable_video_diffusion_pipeline import StableVideoDiffusionPipeline
from utilities import (
PIPELINE_TYPE,
add_arguments,
download_image,
)
from demo_diffusion import dd_argparse
from demo_diffusion import image as image_module
from demo_diffusion import pipeline as pipeline_module
def parseArgs():
parser = argparse.ArgumentParser(description="Options for Stable Diffusion Img2Vid Demo", conflict_handler='resolve')
parser = add_arguments(parser)
parser = dd_argparse.add_arguments(parser)
parser.add_argument('--version', type=str, default="svd-xt-1.1", choices=["svd-xt-1.1"], help="Version of Stable Video Diffusion")
parser.add_argument('--input-image', type=str, default="", help="Path to the input image")
parser.add_argument('--height', type=int, default=576, help="Height of image to generate (must be multiple of 8)")
@@ -44,7 +42,7 @@ def process_pipeline_args(args):
if not args.input_image:
args.input_image = "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/svd/rocket.png?download=true"
if isinstance(args.input_image, str):
input_image = download_image(args.input_image).resize((args.width, args.height))
input_image = image_module.download_image(args.input_image).resize((args.width, args.height))
elif isinstance(args.input_image, Image.Image):
input_image = Image.open(args.input_image)
else:
@@ -61,7 +59,7 @@ def process_pipeline_args(args):
raise ValueError(f"Batch size {args.batch_size} is larger than allowed {max_batch_size}.")
if not args.build_static_batch or args.build_dynamic_shape:
raise ValueError(f"Dynamic shapes not supported. Do not specify `--build-dynamic-shape`")
raise ValueError("Dynamic shapes not supported. Do not specify `--build-dynamic-shape`")
if args.fp8:
import torch
@@ -116,9 +114,9 @@ if __name__ == "__main__":
args = parseArgs()
kwargs_init_pipeline, kwargs_load_engine, args_run_demo = process_pipeline_args(args)
# Initialize demo
demo = StableVideoDiffusionPipeline(
pipeline_type=PIPELINE_TYPE.IMG2VID,
**kwargs_init_pipeline)
demo = pipeline_module.StableVideoDiffusionPipeline(
pipeline_type=pipeline_module.PIPELINE_TYPE.IMG2VID, **kwargs_init_pipeline
)
demo.loadEngines(
args.engine_dir,
args.framework_model_dir,
demo/Diffusion/demo_inpaint.py
+12 -10
View File
@@ -1,5 +1,5 @@
#
# SPDX-FileCopyrightText: Copyright (c) 1993-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-FileCopyrightText: Copyright (c) 1993-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#
# Licensed under the Apache License, Version 2.0 (the "License");
@@ -20,12 +20,14 @@ import argparse
from cuda import cudart
from PIL import Image
from stable_diffusion_pipeline import StableDiffusionPipeline
from utilities import PIPELINE_TYPE, TRT_LOGGER, add_arguments, download_image, process_pipeline_args
from demo_diffusion import dd_argparse
from demo_diffusion import image as image_module
from demo_diffusion import pipeline as pipeline_module
def parseArgs():
parser = argparse.ArgumentParser(description="Options for Stable Diffusion Inpaint Demo", conflict_handler='resolve')
parser = add_arguments(parser)
parser = dd_argparse.add_arguments(parser)
parser.add_argument('--version', type=str, default="1.5", choices=["1.5", "2.0"], help="Stable Diffusion version. Only 1.5 and 2.0 supported for inpainting.")
parser.add_argument('--scheduler', type=str, default="PNDM", choices=["PNDM"], help="Scheduler for diffusion process")
parser.add_argument('--input-image', type=str, default="", help="Path to the input image")
@@ -40,13 +42,13 @@ if __name__ == "__main__":
input_image = Image.open(args.input_image).convert("RGB")
else:
img_url = "https://raw.githubusercontent.com/CompVis/latent-diffusion/main/data/inpainting_examples/overture-creations-5sI6fQgYIuo.png"
input_image = download_image(img_url)
input_image = image_module.download_image(img_url)
if args.mask_image:
mask_image = Image.open(args.mask_image).convert("RGB")
else:
mask_url = "https://raw.githubusercontent.com/CompVis/latent-diffusion/main/data/inpainting_examples/overture-creations-5sI6fQgYIuo_mask.png"
mask_image = download_image(mask_url)
mask_image = image_module.download_image(mask_url)
image_width, image_height = input_image.size
if image_height != args.height or image_width != args.width:
@@ -60,12 +62,12 @@ if __name__ == "__main__":
mask_image = mask_image.resize((args.height, args.width))
mask_height, mask_width = args.height, args.width
kwargs_init_pipeline, kwargs_load_engine, args_run_demo = process_pipeline_args(args)
kwargs_init_pipeline, kwargs_load_engine, args_run_demo = dd_argparse.process_pipeline_args(args)
# Initialize demo
demo = StableDiffusionPipeline(
pipeline_type=PIPELINE_TYPE.INPAINT,
**kwargs_init_pipeline)
demo = pipeline_module.StableDiffusionPipeline(
pipeline_type=pipeline_module.PIPELINE_TYPE.INPAINT, **kwargs_init_pipeline
)
# Load TensorRT engines and pytorch modules
demo.loadEngines(
demo/Diffusion/demo_stable_cascade.py
+14 -12
View File
@@ -1,5 +1,5 @@
#
# SPDX-FileCopyrightText: Copyright (c) 1993-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-FileCopyrightText: Copyright (c) 1993-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#
# Licensed under the Apache License, Version 2.0 (the "License");
@@ -15,18 +15,19 @@
# limitations under the License.
#
import os
import torch
import argparse
import os
import torch
from cuda import cudart
from stable_cascade_pipeline import StableCascadePipeline
from utilities import PIPELINE_TYPE, add_arguments, process_pipeline_args
from demo_diffusion import dd_argparse
from demo_diffusion import pipeline as pipeline_module
def parseArgs():
parser = argparse.ArgumentParser(description="Options for Stable Cascade Txt2Img Demo", conflict_handler='resolve')
parser = add_arguments(parser)
parser = dd_argparse.add_arguments(parser)
parser.add_argument('--version', type=str, default="cascade", choices=["cascade"], help="Version of Stable Cascade")
parser.add_argument('--height', type=int, default=1024, help="Height of image to generate (must be multiple of 8)")
parser.add_argument('--width', type=int, default=1024, help="Width of image to generate (must be multiple of 8)")
@@ -38,19 +39,19 @@ def parseArgs():
return parser.parse_args()
class StableCascadeDemoPipeline(StableCascadePipeline):
class StableCascadeDemoPipeline(pipeline_module.StableCascadePipeline):
def __init__(self, prior_denoising_steps, decoder_denoising_steps, prior_guidance_scale, decoder_guidance_scale, lite, **kwargs):
self.nvtx_profile = kwargs['nvtx_profile']
self.prior = StableCascadePipeline(
pipeline_type=PIPELINE_TYPE.CASCADE_PRIOR,
self.prior = pipeline_module.StableCascadePipeline(
pipeline_type=pipeline_module.PIPELINE_TYPE.CASCADE_PRIOR,
denoising_steps=prior_denoising_steps,
guidance_scale=prior_guidance_scale,
return_latents=True,
lite=lite,
**kwargs,
)
self.decoder = StableCascadePipeline(
pipeline_type=PIPELINE_TYPE.CASCADE_DECODER,
self.decoder = pipeline_module.StableCascadePipeline(
pipeline_type=pipeline_module.PIPELINE_TYPE.CASCADE_DECODER,
denoising_steps=decoder_denoising_steps,
guidance_scale=decoder_guidance_scale,
lite=lite,
@@ -122,11 +123,12 @@ class StableCascadeDemoPipeline(StableCascadePipeline):
self.prior.teardown()
self.decoder.teardown()
if __name__ == "__main__":
print("[I] Initializing StableCascade txt2img demo using TensorRT")
args = parseArgs()
kwargs_init_pipeline, kwargs_load_engine, args_run_demo = process_pipeline_args(args)
kwargs_init_pipeline, kwargs_load_engine, args_run_demo = dd_argparse.process_pipeline_args(args)
# Initialize demo
_ = kwargs_init_pipeline.pop('guidance_scale')
demo/Diffusion/demo_txt2img.py
+9 -8
View File
@@ -1,5 +1,5 @@
#
# SPDX-FileCopyrightText: Copyright (c) 1993-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-FileCopyrightText: Copyright (c) 1993-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#
# Licensed under the Apache License, Version 2.0 (the "License");
@@ -19,24 +19,25 @@ import argparse
from cuda import cudart
from stable_diffusion_pipeline import StableDiffusionPipeline
from utilities import PIPELINE_TYPE, TRT_LOGGER, add_arguments, process_pipeline_args
from demo_diffusion import dd_argparse
from demo_diffusion import pipeline as pipeline_module
def parseArgs():
parser = argparse.ArgumentParser(description="Options for Stable Diffusion Txt2Img Demo")
parser = add_arguments(parser)
parser = dd_argparse.add_arguments(parser)
return parser.parse_args()
if __name__ == "__main__":
print("[I] Initializing StableDiffusion txt2img demo using TensorRT")
args = parseArgs()
kwargs_init_pipeline, kwargs_load_engine, args_run_demo = process_pipeline_args(args)
kwargs_init_pipeline, kwargs_load_engine, args_run_demo = dd_argparse.process_pipeline_args(args)
# Initialize demo
demo = StableDiffusionPipeline(
pipeline_type=PIPELINE_TYPE.TXT2IMG,
**kwargs_init_pipeline)
demo = pipeline_module.StableDiffusionPipeline(
pipeline_type=pipeline_module.PIPELINE_TYPE.TXT2IMG, **kwargs_init_pipeline
)
# Load TensorRT engines and pytorch modules
demo.loadEngines(
demo/Diffusion/demo_txt2img_flux.py
+14 -43
View File
@@ -1,5 +1,5 @@
#
# SPDX-FileCopyrightText: Copyright (c) 1993-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-FileCopyrightText: Copyright (c) 1993-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#
# Licensed under the Apache License, Version 2.0 (the "License");
@@ -16,22 +16,17 @@
#
import argparse
import torch
from cuda import cudart
from flux_pipeline import FluxPipeline
from utilities import (
PIPELINE_TYPE,
add_arguments,
process_pipeline_args,
VALID_OPTIMIZATION_LEVELS,
)
from cuda import cudart
from demo_diffusion import dd_argparse
from demo_diffusion import pipeline as pipeline_module
def parse_args():
parser = argparse.ArgumentParser(
description="Options for Flux Txt2Img Demo", conflict_handler="resolve"
)
parser = add_arguments(parser)
parser = dd_argparse.add_arguments(parser)
parser.add_argument(
"--version",
type=str,
@@ -83,7 +78,7 @@ def parse_args():
"--optimization-level",
type=int,
default=3,
help=f"Set the builder optimization level to build the engine with. A higher level allows TensorRT to spend more building time for more optimization options. Must be one of {VALID_OPTIMIZATION_LEVELS}.",
help=f"Set the builder optimization level to build the engine with. A higher level allows TensorRT to spend more building time for more optimization options. Must be one of {dd_argparse.VALID_OPTIMIZATION_LEVELS}.",
)
parser.add_argument(
"--torch-fallback",
@@ -131,12 +126,6 @@ def parse_args():
result[key] = value
return result
parser.add_argument(
"--model-onnx-dirs",
type=_parse_kv_pairs,
help="Set directories for individual ONNX models. For example: --model-onnx-dirs=transformer:/path/to/transformer,vae:/path/to/vae,t5:/path/to/t5,clip:/path/to/clip"
)
return parser.parse_args()
@@ -174,7 +163,9 @@ def process_demo_args(args):
args.torch_fallback = args.torch_fallback.split(",")
if args.torch_fallback and args.torch_inference:
print(f"[W] All models will run in PyTorch when --torch-inference is set. Parameter --torch-fallback will be ignored.")
print(
"[W] All models will run in PyTorch when --torch-inference is set. Parameter --torch-fallback will be ignored."
)
args.torch_fallback = None
args_run_demo = (
@@ -194,36 +185,17 @@ if __name__ == "__main__":
print("[I] Initializing Flux txt2img demo using TensorRT")
args = parse_args()
kwargs_init_pipeline, kwargs_load_engine, _ = process_pipeline_args(args)
_, kwargs_load_engine, _ = dd_argparse.process_pipeline_args(args)
args_run_demo = process_demo_args(args)
device_info = torch.cuda.get_device_properties(0)
version = device_info.major * 10 + device_info.minor
force_weakly_typed_t5 = True if version == 100 else False
print("[I] Using weakly typed for T5 engine on SM100 if FLUX pipeline is running in fp16")
# Initialize demo
demo = FluxPipeline(
pipeline_type=PIPELINE_TYPE.TXT2IMG,
max_sequence_length=args.max_sequence_length,
bf16=args.bf16,
low_vram=args.low_vram,
torch_fallback=args.torch_fallback,
weight_streaming=args.ws,
t5_weight_streaming_budget_percentage=args.t5_ws_percentage,
transformer_weight_streaming_budget_percentage=args.transformer_ws_percentage,
force_weakly_typed_t5=force_weakly_typed_t5,
**kwargs_init_pipeline)
demo = pipeline_module.FluxPipeline.FromArgs(args, pipeline_type=pipeline_module.PIPELINE_TYPE.TXT2IMG)
# Load TensorRT engines and pytorch modules
demo.load_engines(
args.engine_dir,
args.framework_model_dir,
args.onnx_dir,
framework_model_dir=args.framework_model_dir,
onnx_export_only=args.onnx_export_only,
model_onnx_dirs=args.model_onnx_dirs,
fp4=args.fp4,
**kwargs_load_engine
**kwargs_load_engine,
)
if args.onnx_export_only:
@@ -231,7 +203,6 @@ if __name__ == "__main__":
demo.teardown()
exit(0)
# Since VAE and VAE_encoder require by far the largest device memories, in low-vram mode
# we allocate the required device memory individually before each model is run.
if demo.low_vram:
demo/Diffusion/demo_txt2img_sd3.py
+15 -11
View File
@@ -1,5 +1,5 @@
#
# SPDX-FileCopyrightText: Copyright (c) 1993-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-FileCopyrightText: Copyright (c) 1993-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#
# Licensed under the Apache License, Version 2.0 (the "License");
@@ -16,17 +16,19 @@
#
import argparse
from PIL import Image
from cuda import cudart
from stable_diffusion_3_pipeline import StableDiffusion3Pipeline
from utilities import PIPELINE_TYPE, add_arguments
from utils_sd3.other_impls import preprocess_image_sd3
from cuda import cudart
from PIL import Image
from demo_diffusion import dd_argparse
from demo_diffusion import pipeline as pipeline_module
from demo_diffusion.utils_sd3.other_impls import preprocess_image_sd3
def parseArgs():
# Stable Diffusion 3 configuration
parser = argparse.ArgumentParser(description="Options for Stable Diffusion 3 Txt2Img Demo", conflict_handler='resolve')
parser = add_arguments(parser)
parser = dd_argparse.add_arguments(parser)
parser.add_argument('--version', type=str, default="sd3", choices=["sd3"], help="Version of Stable Diffusion")
parser.add_argument('--height', type=int, default=1024, help="Height of image to generate (must be multiple of 8)")
parser.add_argument('--width', type=int, default=1024, help="Height of image to generate (must be multiple of 8)")
@@ -47,7 +49,9 @@ def process_pipeline_args(args):
raise ValueError(f"Batch size {args.batch_size} is larger than allowed {max_batch_size}.")
if args.use_cuda_graph and (not args.build_static_batch or args.build_dynamic_shape):
raise ValueError(f"Using CUDA graph requires static dimensions. Enable `--build-static-batch` and do not specify `--build-dynamic-shape`")
raise ValueError(
"Using CUDA graph requires static dimensions. Enable `--build-static-batch` and do not specify `--build-dynamic-shape`"
)
input_image = None
if args.input_image:
@@ -100,9 +104,9 @@ if __name__ == "__main__":
kwargs_init_pipeline, kwargs_load_engine, args_run_demo = process_pipeline_args(args)
# Initialize demo
demo = StableDiffusion3Pipeline(
pipeline_type=PIPELINE_TYPE.TXT2IMG,
**kwargs_init_pipeline)
demo = pipeline_module.StableDiffusion3Pipeline(
pipeline_type=pipeline_module.PIPELINE_TYPE.TXT2IMG, **kwargs_init_pipeline
)
# Load TensorRT engines and pytorch modules
demo.loadEngines(
demo/Diffusion/demo_txt2img_xl.py
+17 -12
View File
@@ -1,5 +1,5 @@
#
# SPDX-FileCopyrightText: Copyright (c) 1993-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-FileCopyrightText: Copyright (c) 1993-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#
# Licensed under the Apache License, Version 2.0 (the "License");
@@ -19,12 +19,13 @@ import argparse
from cuda import cudart
from stable_diffusion_pipeline import StableDiffusionPipeline
from utilities import PIPELINE_TYPE, TRT_LOGGER, add_arguments, process_pipeline_args
from demo_diffusion import dd_argparse
from demo_diffusion import pipeline as pipeline_module
def parseArgs():
parser = argparse.ArgumentParser(description="Options for Stable Diffusion XL Txt2Img Demo", conflict_handler='resolve')
parser = add_arguments(parser)
parser = dd_argparse.add_arguments(parser)
parser.add_argument('--version', type=str, default="xl-1.0", choices=["xl-1.0", "xl-turbo"], help="Version of Stable Diffusion XL")
parser.add_argument('--height', type=int, default=1024, help="Height of image to generate (must be multiple of 8)")
parser.add_argument('--width', type=int, default=1024, help="Height of image to generate (must be multiple of 8)")
@@ -39,21 +40,24 @@ def parseArgs():
return parser.parse_args()
class StableDiffusionXLPipeline(StableDiffusionPipeline):
class StableDiffusionXLPipeline(pipeline_module.StableDiffusionPipeline):
def __init__(self, vae_scaling_factor=0.13025, enable_refiner=False, **kwargs):
self.enable_refiner = enable_refiner
self.nvtx_profile = kwargs['nvtx_profile']
self.base = StableDiffusionPipeline(
pipeline_type=PIPELINE_TYPE.XL_BASE,
self.base = pipeline_module.StableDiffusionPipeline(
pipeline_type=pipeline_module.PIPELINE_TYPE.XL_BASE,
vae_scaling_factor=vae_scaling_factor,
return_latents=self.enable_refiner,
**kwargs)
**kwargs,
)
if self.enable_refiner:
self.refiner = StableDiffusionPipeline(
pipeline_type=PIPELINE_TYPE.XL_REFINER,
self.refiner = pipeline_module.StableDiffusionPipeline(
pipeline_type=pipeline_module.PIPELINE_TYPE.XL_REFINER,
vae_scaling_factor=vae_scaling_factor,
return_latents=False,
**kwargs)
**kwargs,
)
def loadEngines(self, framework_model_dir, onnx_dir, engine_dir, onnx_refiner_dir='onnx_xl_refiner', engine_refiner_dir='engine_xl_refiner', **kwargs):
self.base.loadEngines(engine_dir, framework_model_dir, onnx_dir, **kwargs)
@@ -121,11 +125,12 @@ class StableDiffusionXLPipeline(StableDiffusionPipeline):
if self.enable_refiner:
self.refiner.teardown()
if __name__ == "__main__":
print("[I] Initializing TensorRT accelerated StableDiffusionXL txt2img pipeline")
args = parseArgs()
kwargs_init_pipeline, kwargs_load_engine, args_run_demo = process_pipeline_args(args)
kwargs_init_pipeline, kwargs_load_engine, args_run_demo = dd_argparse.process_pipeline_args(args)
# Initialize demo
demo = StableDiffusionXLPipeline(vae_scaling_factor=0.13025, enable_refiner=args.enable_refiner, **kwargs_init_pipeline)
demo/Diffusion/docs/support_matrix.md
+33
View File
@@ -0,0 +1,33 @@
<!--- SPDX-License-Identifier: Apache-2.0 -->
# Supported Diffusion Models
This demo supports Diffusion models that are popular in the Generative AI community. The table below lists the various configurations we support for each pipeline.
## Pipeline Support Matrix
| Pipeline | Version | Task | Supported Precisions | Hub | Restrictions |
|------------------------|------------------------------------------------------------------------------------------------------------------|-------------------------------------------------------------|----------------------|----------------------------------|-----------------------------------------------------------------------------------------------------------------------------------------------------------------------|
| Stable Diffusion | 1.4 | <ul><li>Text-to-image</li><li>Image-to-image</li></ul> | FP16, FP8, INT8 | [CompVis/stable-diffusion-v1-4](https://huggingface.co/CompVis/stable-diffusion-v1-4) |
| Stable Diffusion | [1.5](../README.md#generate-an-image-guided-by-a-text-prompt) | <ul><li>Text-to-image</li><li>Image-to-image</li></ul> | FP16, FP8, INT8 | [KiwiXR/stable-diffusion-v1-5](https://huggingface.co/KiwiXR/stable-diffusion-v1-5) |
| Stable Diffusion | 1.4, [1.5](../README.md#generate-an-inpainted-image-guided-by-an-image-mask-and-a-text-prompt) | <ul><li>Inpainting</li></ul> | FP16 | [benjamin-paine/stable-diffusion-v1-5-inpainting](https://huggingface.co/benjamin-paine/stable-diffusion-v1-5-inpainting) |
| Stable Diffusion | dreamshaper-7 | <ul><li>Text-to-image</li><li>Image-to-image</li></ul> | FP16 | [Lykon/dreamshaper-7](https://huggingface.co/Lykon/dreamshaper-7) |
| Stable Diffusion | 2.0-base | <ul><li>Text-to-image</li><li>Image-to-image</li></ul> | FP16 | [stabilityai/stable-diffusion-2-base](https://huggingface.co/stabilityai/stable-diffusion-2-base) |
| Stable Diffusion | 2.0 | <ul><li>Text-to-image</li><li>Image-to-image</li></ul> | FP16 | [stabilityai/stable-diffusion-2](https://huggingface.co/stabilityai/stable-diffusion-2) |
| Stable Diffusion | 2.0, 2.0-base | <ul><li>Inpainting</li></ul> | FP16 | [stabilityai/stable-diffusion-2-inpainting](https://huggingface.co/stabilityai/stable-diffusion-2-inpainting) |
| Stable Diffusion | 2.1-base | <ul><li>Text-to-image</li><li>Image-to-image</li></ul> | FP16, FP8, INT8 | [stabilityai/stable-diffusion-2-1-base](https://huggingface.co/stabilityai/stable-diffusion-2-1-base) |
| Stable Diffusion | 2.1 | <ul><li>Text-to-image</li><li>Image-to-image</li></ul> | FP16, FP8, INT8 | [stabilityai/stable-diffusion-2-1](https://huggingface.co/stabilityai/stable-diffusion-2-1) |
| Stable Diffusion | [XL 1.0-base](../README.md#generate-an-image-with-stable-diffusion-xl-guided-by-a-single-text-prompt) | <ul><li>Text-to-image</li><li>Image-to-image</li></ul> | FP16, FP8, INT8 | [stabilityai/stable-diffusion-xl-base-1.0](https://huggingface.co/stabilityai/stable-diffusion-xl-base-1.0) |
| Stable Diffusion | [XL 1.0-refiner](../README.md#generate-an-image-with-stable-diffusion-xl-guided-by-a-single-text-prompt) | <ul><li>Text-to-image</li><li>Image-to-image</li></ul> | FP16 | [stabilityai/stable-diffusion-xl-refiner-1.0](https://huggingface.co/stabilityai/stable-diffusion-xl-refiner-1.0) |
| Stable Diffusion | [XL-Turbo](../README.md#faster-text-to-image-using-sdxl-turbo) | <ul><li>Text-to-image</li><li>Image-to-image</li></ul> | FP16 | [stabilityai/sdxl-turbo](https://huggingface.co/stabilityai/sdxl-turbo) |
| Stable Diffusion | [3](../README.md#generate-an-image-guided-by-a-text-prompt-using-stable-diffusion-3) | <ul><li>Text-to-image</li></ul> | FP16 | [stabilityai/stable-diffusion-3-medium](https://huggingface.co/stabilityai/stable-diffusion-3-medium) |
| ControlNet | [1.5](../README.md#generate-an-image-with-controlnet-guided-by-images-and-text-prompts) | <ul><li>Image-to-image</li></ul> | FP16 | <ul><li>[lllyasviel/sd-controlnet-canny](https://huggingface.co/lllyasviel/sd-controlnet-canny)</li><li>[lllyasviel/sd-controlnet-depth](https://huggingface.co/lllyasviel/sd-controlnet-depth)</li><li>[lllyasviel/sd-controlnet-hed](https://huggingface.co/lllyasviel/sd-controlnet-hed)</li><li>[lllyasviel/sd-controlnet-mlsd](https://huggingface.co/lllyasviel/sd-controlnet-mlsd)</li><li>[lllyasviel/sd-controlnet-normal](https://huggingface.co/lllyasviel/sd-controlnet-normal)</li><li>[lllyasviel/sd-controlnet_openpose](https://huggingface.co/lllyasviel/sd-controlnet-openpose)</li><li>[lllyasviel/sd-controlnet_scribble](https://huggingface.co/lllyasviel/sd-controlnet-scribble)</li><li>[lllyasviel/sd-controlnet_seg](https://huggingface.co/lllyasviel/sd-controlnet-seg)</li></ul> |
| ControlNet | XL 1.0-base | <ul><li>Image-to-image</li></ul> | FP16 | [diffusers/controlnet-canny-sdxl-1.0](https://huggingface.co/diffusers/controlnet-canny-sdxl-1.0) |
| Stable Video Diffusion | [XT-1.1](../README.md#generate-a-video-guided-by-an-initial-image-using-stable-video-diffusion) | <ul><li>Text-to-video</li></ul> | FP16, FP8 | [stabilityai/stable-video-diffusion-img2vid-xt-1-1](https://huggingface.co/stabilityai/stable-video-diffusion-img2vid-xt-1-1) |
| Stable Cascade | [N/A](../README.md#generate-an-image-guided-by-a-text-prompt-using-stable-cascade) | <ul><li>Text-to-image</li></ul> | BF16 | <ul><li>[stabilityai/stable-cascade-prior](https://huggingface.co/stabilityai/stable-cascade-prior)</li><li>[stabilityai/stable-cascade](https://huggingface.co/stabilityai/stable-cascade)</li></ul> |
| Flux | [1-Dev](../README.md#generate-an-image-guided-by-a-text-prompt-using-flux) | <ul><li>Text-to-image</li><li>Image-to-image</li></ul> | FP16, BF16, FP8, FP4 | [black-forest-labs/FLUX.1-dev](https://huggingface.co/black-forest-labs/FLUX.1-dev) |
| Flux | [1-Schnell](../README.md#generate-an-image-guided-by-a-text-prompt-using-flux) | <ul><li>Text-to-image</li><li>Image-to-image</li></ul> | FP16, BF16, FP8, FP4 | [black-forest-labs/FLUX.1-schnell](https://huggingface.co/black-forest-labs/FLUX.1-schnell) |
| Flux | [1-Canny-Dev](../README.md#generate-an-image-guided-by-a-text-prompt-and-a-control-image-using-flux-controlnet) | <ul><li>Image-to-image</li></ul> | FP16, BF16, FP8, FP4 | [black-forest-labs/FLUX.1-Canny-dev](https://huggingface.co/black-forest-labs/FLUX.1-Canny-dev) |
| Flux | [1-Depth-Dev](../README.md#generate-an-image-guided-by-a-text-prompt-and-a-control-image-using-flux-controlnet) | <ul><li>Image-to-image</li></ul> | FP16, BF16, FP8, FP4 | [black-forest-labs/FLUX.1-Depth-dev](https://huggingface.co/black-forest-labs/FLUX.1-Depth-dev) |
NOTE: Only the text2image pipelines support FP8 and INT8 quantization. The image2image pipelines don't support quantization.
demo/Diffusion/models.py
-2220
View File
File diff suppressed because it is too large Load Diff
demo/Diffusion/scripts/download_flux_onnx_models.sh
-142
View File
@@ -1,142 +0,0 @@
# Initialize Git LFS
git lfs install
export GIT_LFS_SKIP_SMUDGE=1
# Set argument defaults
arg_version="flux.1-dev"
arg_precision="bf16"
arg_help=0
while [[ "$#" -gt 0 ]]; do case $1 in
--version) arg_version="$2"; shift;;
--precision) arg_precision="$2"; shift;;
-h|--help) arg_help=1;;
*) echo "Unknown parameter passed: $1"; echo "For help type: $0 --help"; exit 1;
esac; shift; done
if [ "$arg_help" -eq "1" ]; then
echo "Usage: $0 [options]"
echo " --help or -h : Print this help menu."
echo " --version <Flux Version> : Choose one of ["flux.1-dev", "flux.1-schnell", "flux.1-dev-canny", "flux.1-dev-depth"]"
echo " --precision <Model Precision> : Choose one of ["bf16", "fp8", "fp4"]"
exit;
fi
if [ "$arg_version" = "flux.1-dev" ]; then
# Clone the repository if it doesn't already exist
onnx_dir="onnx-flux-dev"
if [ ! -d "$onnx_dir" ] ; then
git clone https://huggingface.co/black-forest-labs/FLUX.1-dev-onnx $onnx_dir
fi
cd $onnx_dir
git lfs pull --include=clip.opt
git lfs pull --include=t5.opt
git lfs pull --include=vae.opt
cd transformer.opt
if [ "$arg_precision" = "bf16" ]; then
git lfs pull --include=bf16/
mkdir -p ../transformer_bf16/transformer.opt
ln -s $PWD/bf16/* $PWD/../transformer_bf16/transformer.opt
elif [ "$arg_precision" = "fp8" ]; then
git lfs pull --include=fp8/
mkdir -p ../transformer_fp8/transformer-fp8.l4.0.bs2.s50.c32.p1.0.a0.8.opt/
ln -s $PWD/fp8/* $PWD/../transformer_fp8/transformer-fp8.l4.0.bs2.s50.c32.p1.0.a0.8.opt/
elif [ "$arg_precision" = "fp4" ]; then
git lfs pull --include=fp4/
mkdir -p ../transformer_fp4/transformer.opt
ln -s $PWD/fp4/* $PWD/../transformer_fp4/transformer.opt
else
echo "Precision input $arg_precision not supported. Please choose one of ["bf16", "fp8", "fp4"]"
cd ../..
fi
elif [ "$arg_version" = "flux.1-schnell" ]; then
# Clone the repository if it doesn't already exist
onnx_dir="onnx-flux-schnell"
if [ ! -d "$onnx_dir" ] ; then
git clone https://huggingface.co/black-forest-labs/FLUX.1-schnell-onnx $onnx_dir
fi
cd $onnx_dir
git lfs pull --include=clip.opt
git lfs pull --include=t5.opt
git lfs pull --include=vae.opt
cd transformer.opt
if [ "$arg_precision" = "bf16" ]; then
git lfs pull --include=bf16/
mkdir -p ../transformer_bf16/transformer.opt
ln -s $PWD/bf16/* $PWD/../transformer_bf16/transformer.opt
elif [ "$arg_precision" = "fp8" ]; then
git lfs pull --include=fp8/
mkdir -p ../transformer_fp8/transformer-fp8.l4.0.bs2.s50.c32.p1.0.a0.8.opt/
ln -s $PWD/fp8/* $PWD/../transformer_fp8/transformer-fp8.l4.0.bs2.s50.c32.p1.0.a0.8.opt/
elif [ "$arg_precision" = "fp4" ]; then
git lfs pull --include=fp4/
mkdir -p ../transformer_fp4/transformer.opt
ln -s $PWD/fp4/* $PWD/../transformer_fp4/transformer.opt
else
echo "Precision input $arg_precision not supported. Please choose one of ["bf16", "fp8", "fp4"]"
cd ../..
fi
elif [ "$arg_version" = "flux.1-dev-depth" ]; then
# Clone the repository if it doesn't already exist
onnx_dir="onnx-flux-dev-depth"
if [ ! -d "$onnx_dir" ] ; then
git clone https://huggingface.co/black-forest-labs/FLUX.1-Depth-dev-onnx $onnx_dir
fi
cd $onnx_dir
git lfs pull --include=clip.opt
git lfs pull --include=t5.opt
git lfs pull --include=vae.opt
git lfs pull --include=vae_encoder.opt
cd transformer.opt
if [ "$arg_precision" = "bf16" ]; then
git lfs pull --include=bf16/
mkdir -p ../transformer_bf16/transformer.opt
ln -s $PWD/bf16/* $PWD/../transformer_bf16/transformer.opt
elif [ "$arg_precision" = "fp8" ]; then
git lfs pull --include=fp8/
mkdir -p ../transformer_fp8/transformer-fp8.l4.0.bs2.s30.c32.p1.0.a0.8.opt/
ln -s $PWD/fp8/* $PWD/../transformer_fp8/transformer-fp8.l4.0.bs2.s30.c32.p1.0.a0.8.opt/
elif [ "$arg_precision" = "fp4" ]; then
git lfs pull --include=fp4/
mkdir -p ../transformer_fp4/transformer.opt
ln -s $PWD/fp4/* $PWD/../transformer_fp4/transformer.opt
else
echo "Precision input $arg_precision not supported. Please choose one of ["bf16", "fp8", "fp4"]"
cd ../..
fi
elif [ "$arg_version" = "flux.1-dev-canny" ]; then
# Clone the repository if it doesn't already exist
onnx_dir="onnx-flux-dev-canny"
if [ ! -d "$onnx_dir" ] ; then
git clone https://huggingface.co/black-forest-labs/FLUX.1-Canny-dev-onnx $onnx_dir
fi
cd $onnx_dir
git lfs pull --include=clip.opt
git lfs pull --include=t5.opt
git lfs pull --include=vae.opt
git lfs pull --include=vae_encoder.opt
cd transformer.opt
if [ "$arg_precision" = "bf16" ]; then
git lfs pull --include=bf16/
mkdir -p ../transformer_bf16/transformer.opt
ln -s $PWD/bf16/* $PWD/../transformer_bf16/transformer.opt
elif [ "$arg_precision" = "fp8" ]; then
git lfs pull --include=fp8/
mkdir -p ../transformer_fp8/transformer-fp8.l4.0.bs2.s30.c32.p1.0.a0.8.opt/
ln -s $PWD/fp8/* $PWD/../transformer_fp8/transformer-fp8.l4.0.bs2.s30.c32.p1.0.a0.8.opt/
elif [ "$arg_precision" = "fp4" ]; then
git lfs pull --include=fp4/
mkdir -p ../transformer_fp4/transformer.opt
ln -s $PWD/fp4/* $PWD/../transformer_fp4/transformer.opt
else
echo "Precision input $arg_precision not supported. Please choose one of ["bf16", "fp8", "fp4"]"
cd ../..
fi
else
echo "Version $arg_version not supported. Please choose one of ["flux.1-dev", "flux.1-schnell", "flux.1-dev-canny", "flux.1-dev-depth"]"
fi
demo/Diffusion/utilities.py
-806
View File
@@ -1,806 +0,0 @@
#
# Copyright (c) Alibaba, Inc. and its affiliates.
# SPDX-FileCopyrightText: Copyright (c) 1993-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
import gc
import os
import random
import types
import warnings
from collections import OrderedDict
from enum import Enum, auto
from importlib import import_module
from io import BytesIO
import numpy as np
import onnx
import requests
import tensorrt as trt
import torch
from cuda import cudart
from diffusers.models.lora import LoRACompatibleConv, LoRACompatibleLinear
from diffusers.utils import load_image
from onnx import numpy_helper
from PIL import Image
from polygraphy.backend.common import bytes_from_path
from polygraphy.backend.trt import (
CreateConfig,
ModifyNetworkOutputs,
Profile,
engine_from_bytes,
engine_from_network,
network_from_onnx_path,
save_engine,
)
from polygraphy.logger import G_LOGGER
TRT_LOGGER = trt.Logger(trt.Logger.ERROR)
def GiB(val):
return val * 1 << 30
# Map of TensorRT dtype -> torch dtype
trt_to_torch_dtype_dict = {
trt.DataType.BOOL : torch.bool,
trt.DataType.UINT8 : torch.uint8,
trt.DataType.INT8 : torch.int8,
trt.DataType.INT32 : torch.int32,
trt.DataType.INT64 : torch.int64,
trt.DataType.HALF : torch.float16,
trt.DataType.FLOAT : torch.float32,
trt.DataType.BF16 : torch.bfloat16
}
# Define valid optimization levels for TensorRT engine build
VALID_OPTIMIZATION_LEVELS = list(range(6))
def import_from_diffusers(model_name, module_name):
try:
module = import_module(module_name)
return getattr(module, model_name)
except ImportError:
warnings.warn(f"Failed to import {module_name}. The {model_name} model will not be available.", ImportWarning)
except AttributeError:
warnings.warn(f"The {model_name} model is not available in the installed version of diffusers.", ImportWarning)
return None
def unload_model(model):
if model:
del model
torch.cuda.empty_cache()
gc.collect()
def replace_lora_layers(model):
def lora_forward(self, x, scale=None):
return self._torch_forward(x)
for name, module in model.named_modules():
if isinstance(module, LoRACompatibleConv):
in_channels = module.in_channels
out_channels = module.out_channels
kernel_size = module.kernel_size
stride = module.stride
padding = module.padding
dilation = module.dilation
groups = module.groups
bias = module.bias
new_conv = torch.nn.Conv2d(
in_channels,
out_channels,
kernel_size,
stride=stride,
padding=padding,
dilation=dilation,
groups=groups,
bias=bias is not None,
)
new_conv.weight.data = module.weight.data.clone().to(module.weight.data.device)
if bias is not None:
new_conv.bias.data = module.bias.data.clone().to(module.bias.data.device)
# Replace the LoRACompatibleConv layer with the Conv2d layer
path = name.split(".")
sub_module = model
for p in path[:-1]:
sub_module = getattr(sub_module, p)
setattr(sub_module, path[-1], new_conv)
new_conv._torch_forward = new_conv.forward
new_conv.forward = types.MethodType(lora_forward, new_conv)
elif isinstance(module, LoRACompatibleLinear):
in_features = module.in_features
out_features = module.out_features
bias = module.bias
new_linear = torch.nn.Linear(in_features, out_features, bias=bias is not None)
new_linear.weight.data = module.weight.data.clone().to(module.weight.data.device)
if bias is not None:
new_linear.bias.data = module.bias.data.clone().to(module.bias.data.device)
# Replace the LoRACompatibleLinear layer with the Linear layer
path = name.split(".")
sub_module = model
for p in path[:-1]:
sub_module = getattr(sub_module, p)
setattr(sub_module, path[-1], new_linear)
new_linear._torch_forward = new_linear.forward
new_linear.forward = types.MethodType(lora_forward, new_linear)
def merge_loras(model, lora_loader):
paths, weights, scale = lora_loader.paths, lora_loader.weights, lora_loader.scale
for i, path in enumerate(paths):
print(f"[I] Loading LoRA: {path}, weight {weights[i]}")
state_dict, network_alphas = lora_loader.lora_state_dict(path, unet_config=model.config)
lora_loader.load_lora_into_unet(state_dict, network_alphas=network_alphas,
unet=model, adapter_name=path)
model.set_adapters(paths, weights=weights)
# NOTE: fuse_lora an experimental API in Diffusers
model.fuse_lora(adapter_names=paths, lora_scale=scale)
model.unload_lora()
return model
def CUASSERT(cuda_ret):
err = cuda_ret[0]
if err != cudart.cudaError_t.cudaSuccess:
raise RuntimeError(f"CUDA ERROR: {err}, error code reference: https://nvidia.github.io/cuda-python/module/cudart.html#cuda.cudart.cudaError_t")
if len(cuda_ret) > 1:
return cuda_ret[1]
return None
class PIPELINE_TYPE(Enum):
TXT2IMG = auto()
IMG2IMG = auto()
IMG2VID = auto()
INPAINT = auto()
CONTROLNET = auto()
XL_BASE = auto()
XL_REFINER = auto()
CASCADE_PRIOR = auto()
CASCADE_DECODER = auto()
def is_txt2img(self):
return self == self.TXT2IMG
def is_img2img(self):
return self == self.IMG2IMG
def is_img2vid(self):
return self == self.IMG2VID
def is_inpaint(self):
return self == self.INPAINT
def is_controlnet(self):
return self == self.CONTROLNET
def is_sd_xl_base(self):
return self == self.XL_BASE
def is_sd_xl_refiner(self):
return self == self.XL_REFINER
def is_sd_xl(self):
return self.is_sd_xl_base() or self.is_sd_xl_refiner()
def is_cascade_prior(self):
return self == self.CASCADE_PRIOR
def is_cascade_decoder(self):
return self == self.CASCADE_DECODER
def is_cascade(self):
return self.is_cascade_prior() or self.is_cascade_decoder()
class Engine():
def __init__(
self,
engine_path,
):
self.engine_path = engine_path
self.engine = None
self.context = None
self.buffers = OrderedDict()
self.tensors = OrderedDict()
self.cuda_graph_instance = None # cuda graph
def __del__(self):
del self.engine
del self.context
del self.buffers
del self.tensors
def refit(self, refit_weights, updated_weight_names):
# Initialize refitter
refitter = trt.Refitter(self.engine, TRT_LOGGER)
refitted_weights = set()
def refit_single_weight(trt_weight_name):
# get weight from state dict
trt_datatype = refitter.get_weights_prototype(trt_weight_name).dtype
refit_weights[trt_weight_name] = refit_weights[trt_weight_name].to(trt_to_torch_dtype_dict[trt_datatype])
# trt.Weight and trt.TensorLocation
trt_wt_tensor = trt.Weights(trt_datatype, refit_weights[trt_weight_name].data_ptr(), torch.numel(refit_weights[trt_weight_name]))
trt_wt_location = trt.TensorLocation.DEVICE if refit_weights[trt_weight_name].is_cuda else trt.TensorLocation.HOST
# apply refit
refitter.set_named_weights(trt_weight_name, trt_wt_tensor, trt_wt_location)
refitted_weights.add(trt_weight_name)
# iterate through all tensorrt refittable weights
for trt_weight_name in refitter.get_all_weights():
if trt_weight_name not in updated_weight_names:
continue
refit_single_weight(trt_weight_name)
# iterate through missing weights required by tensorrt - addresses the case where lora_scale=0
for trt_weight_name in refitter.get_missing_weights():
refit_single_weight(trt_weight_name)
if not refitter.refit_cuda_engine():
print("Error: failed to refit new weights.")
exit(0)
print(f"[I] Total refitted weights {len(refitted_weights)}.")
def build(self,
onnx_path,
strongly_typed=False,
fp16=True,
bf16=False,
tf32=False,
int8=False,
fp8=False,
input_profile=None,
enable_refit=False,
enable_all_tactics=False,
timing_cache=None,
update_output_names=None,
native_instancenorm=True,
verbose=False,
weight_streaming=False,
**extra_build_args
):
print(f"Building TensorRT engine for {onnx_path}: {self.engine_path}")
p = Profile()
if input_profile:
for name, dims in input_profile.items():
assert len(dims) == 3
p.add(name, min=dims[0], opt=dims[1], max=dims[2])
if not enable_all_tactics:
extra_build_args['tactic_sources'] = []
flags = []
if native_instancenorm:
flags.append(trt.OnnxParserFlag.NATIVE_INSTANCENORM)
# Weight streaming requires the engine to have strong typing, therefore builder flags specifying precision, such as int8 and fp16, should not be enabled.
# Please find more details in our developer guide: https://docs.nvidia.com/deeplearning/tensorrt/developer-guide/index.html#streaming-weights.
if weight_streaming:
strongly_typed = True
fp16 = False
bf16 = False
int8 = False
fp8 = False
print(f"Strongly typed mode is {strongly_typed} for {onnx_path}")
network = network_from_onnx_path(
onnx_path,
flags=flags,
strongly_typed=strongly_typed
)
if update_output_names:
print(f"Updating network outputs to {update_output_names}")
network = ModifyNetworkOutputs(network, update_output_names)
with G_LOGGER.verbosity(G_LOGGER.EXTRA_VERBOSE if verbose else G_LOGGER.ERROR):
engine = engine_from_network(
network,
config=CreateConfig(fp16=fp16,
bf16=bf16,
tf32=tf32,
int8=int8,
fp8=fp8,
refittable=enable_refit,
profiles=[p],
load_timing_cache=timing_cache,
weight_streaming=weight_streaming,
**extra_build_args
),
save_timing_cache=timing_cache
)
save_engine(engine, path=self.engine_path)
def load(self, weight_streaming=False, weight_streaming_budget_percentage=None):
if self.engine is not None:
print(f"[W]: Engine {self.engine_path} already loaded, skip reloading")
return
if not hasattr(self,'engine_bytes_cpu') or self.engine_bytes_cpu is None:
# keep a cpu copy of the engine to reduce reloading time.
print(f"Loading TensorRT engine to cpu bytes: {self.engine_path}")
self.engine_bytes_cpu = bytes_from_path(self.engine_path)
print(f"Loading TensorRT engine from bytes: {self.engine_path}")
self.engine = engine_from_bytes(self.engine_bytes_cpu)
if weight_streaming:
if weight_streaming_budget_percentage is None:
warnings.warn(f"Weight streaming budget is not set for {self.engine_path}. Weights will not be streamed.")
else:
self.engine.weight_streaming_budget_v2 = int(weight_streaming_budget_percentage / 100 * self.engine.streamable_weights_size)
def unload(self):
if self.engine is not None:
print(f"Unloading TensorRT engine: {self.engine_path}")
del self.engine
self.engine = None
gc.collect()
else:
print(f"[W]: Unload an unloaded engine {self.engine_path}, skip unloading")
def activate(self, device_memory=None):
if device_memory:
self.context = self.engine.create_execution_context_without_device_memory()
self.context.device_memory = device_memory
else:
self.context = self.engine.create_execution_context()
def reactivate(self, device_memory):
assert self.context
self.context.device_memory = device_memory
def deactivate(self):
del self.context
self.context = None
def allocate_buffers(self, shape_dict=None, device='cuda'):
for binding in range(self.engine.num_io_tensors):
name = self.engine.get_tensor_name(binding)
if shape_dict and name in shape_dict:
shape = shape_dict[name]
else:
shape = self.engine.get_tensor_shape(name)
print(f"[W]: {self.engine_path}: Could not find '{name}' in shape dict {shape_dict}. Using shape {shape} inferred from the engine.")
if self.engine.get_tensor_mode(name) == trt.TensorIOMode.INPUT:
self.context.set_input_shape(name, shape)
dtype=trt_to_torch_dtype_dict[self.engine.get_tensor_dtype(name)]
tensor = torch.empty(tuple(shape), dtype=dtype).to(device=device)
self.tensors[name] = tensor
def deallocate_buffers(self):
for idx in range(self.engine.num_io_tensors):
binding = self.engine[idx]
del self.tensors[binding]
def infer(self, feed_dict, stream, use_cuda_graph=False):
for name, buf in feed_dict.items():
self.tensors[name].copy_(buf)
for name, tensor in self.tensors.items():
self.context.set_tensor_address(name, tensor.data_ptr())
if use_cuda_graph:
if self.cuda_graph_instance is not None:
CUASSERT(cudart.cudaGraphLaunch(self.cuda_graph_instance, stream))
CUASSERT(cudart.cudaStreamSynchronize(stream))
else:
# do inference before CUDA graph capture
noerror = self.context.execute_async_v3(stream)
if not noerror:
raise ValueError("ERROR: inference failed.")
# capture cuda graph
CUASSERT(cudart.cudaStreamBeginCapture(stream, cudart.cudaStreamCaptureMode.cudaStreamCaptureModeGlobal))
self.context.execute_async_v3(stream)
self.graph = CUASSERT(cudart.cudaStreamEndCapture(stream))
self.cuda_graph_instance = CUASSERT(cudart.cudaGraphInstantiate(self.graph, 0))
else:
noerror = self.context.execute_async_v3(stream)
if not noerror:
raise ValueError("ERROR: inference failed.")
return self.tensors
def save_image(images, image_path_dir, image_name_prefix, image_name_suffix):
"""
Save the generated images to png files.
"""
for i in range(images.shape[0]):
image_path = os.path.join(image_path_dir, image_name_prefix+str(i+1)+'-'+str(random.randint(1000,9999))+'-'+image_name_suffix+'.png')
print(f"Saving image {i+1} / {images.shape[0]} to: {image_path}")
Image.fromarray(images[i]).save(image_path)
def preprocess_image(image):
"""
image: torch.Tensor
"""
w, h = image.size
w, h = map(lambda x: x - x % 32, (w, h)) # resize to integer multiple of 32
image = image.resize((w, h))
image = np.array(image).astype(np.float32) / 255.0
image = image[None].transpose(0, 3, 1, 2)
image = torch.from_numpy(image).contiguous()
return 2.0 * image - 1.0
# Taken from https://github.com/huggingface/diffusers/blob/main/src/diffusers/pipelines/stable_video_diffusion/pipeline_stable_video_diffusion.py#L620
def _resize_with_antialiasing(input, size, interpolation="bicubic", align_corners=True):
h, w = input.shape[-2:]
factors = (h / size[0], w / size[1])
# First, we have to determine sigma
# Taken from skimage: https://github.com/scikit-image/scikit-image/blob/v0.19.2/skimage/transform/_warps.py#L171
sigmas = (
max((factors[0] - 1.0) / 2.0, 0.001),
max((factors[1] - 1.0) / 2.0, 0.001),
)
# Now kernel size. Good results are for 3 sigma, but that is kind of slow. Pillow uses 1 sigma
# https://github.com/python-pillow/Pillow/blob/master/src/libImaging/Resample.c#L206
# But they do it in the 2 passes, which gives better results. Let's try 2 sigmas for now
ks = int(max(2.0 * 2 * sigmas[0], 3)), int(max(2.0 * 2 * sigmas[1], 3))
# Make sure it is odd
if (ks[0] % 2) == 0:
ks = ks[0] + 1, ks[1]
if (ks[1] % 2) == 0:
ks = ks[0], ks[1] + 1
input = _gaussian_blur2d(input, ks, sigmas)
output = torch.nn.functional.interpolate(input, size=size, mode=interpolation, align_corners=align_corners)
return output
def _compute_padding(kernel_size):
"""Compute padding tuple."""
# 4 or 6 ints: (padding_left, padding_right,padding_top,padding_bottom)
# https://pytorch.org/docs/stable/nn.html#torch.nn.functional.pad
if len(kernel_size) < 2:
raise AssertionError(kernel_size)
computed = [k - 1 for k in kernel_size]
# for even kernels we need to do asymmetric padding :(
out_padding = 2 * len(kernel_size) * [0]
for i in range(len(kernel_size)):
computed_tmp = computed[-(i + 1)]
pad_front = computed_tmp // 2
pad_rear = computed_tmp - pad_front
out_padding[2 * i + 0] = pad_front
out_padding[2 * i + 1] = pad_rear
return out_padding
def _filter2d(input, kernel):
# prepare kernel
b, c, h, w = input.shape
tmp_kernel = kernel[:, None, ...].to(device=input.device, dtype=input.dtype)
tmp_kernel = tmp_kernel.expand(-1, c, -1, -1)
height, width = tmp_kernel.shape[-2:]
padding_shape: list[int] = _compute_padding([height, width])
input = torch.nn.functional.pad(input, padding_shape, mode="reflect")
# kernel and input tensor reshape to align element-wise or batch-wise params
tmp_kernel = tmp_kernel.reshape(-1, 1, height, width)
input = input.view(-1, tmp_kernel.size(0), input.size(-2), input.size(-1))
# convolve the tensor with the kernel.
output = torch.nn.functional.conv2d(input, tmp_kernel, groups=tmp_kernel.size(0), padding=0, stride=1)
out = output.view(b, c, h, w)
return out
def _gaussian(window_size: int, sigma):
if isinstance(sigma, float):
sigma = torch.tensor([[sigma]])
batch_size = sigma.shape[0]
x = (torch.arange(window_size, device=sigma.device, dtype=sigma.dtype) - window_size // 2).expand(batch_size, -1)
if window_size % 2 == 0:
x = x + 0.5
gauss = torch.exp(-x.pow(2.0) / (2 * sigma.pow(2.0)))
return gauss / gauss.sum(-1, keepdim=True)
def _gaussian_blur2d(input, kernel_size, sigma):
if isinstance(sigma, tuple):
sigma = torch.tensor([sigma], dtype=input.dtype)
else:
sigma = sigma.to(dtype=input.dtype)
ky, kx = int(kernel_size[0]), int(kernel_size[1])
bs = sigma.shape[0]
kernel_x = _gaussian(kx, sigma[:, 1].view(bs, 1))
kernel_y = _gaussian(ky, sigma[:, 0].view(bs, 1))
out_x = _filter2d(input, kernel_x[..., None, :])
out = _filter2d(out_x, kernel_y[..., None])
return out
def _append_dims(x, target_dims):
"""Appends dimensions to the end of a tensor until it has target_dims dimensions."""
dims_to_append = target_dims - x.ndim
if dims_to_append < 0:
raise ValueError(f"input has {x.ndim} dims but target_dims is {target_dims}, which is less")
return x[(...,) + (None,) * dims_to_append]
# Not a contribution
# Changes made by NVIDIA CORPORATION & AFFILIATES enabling tensor2vid or otherwise documented as
# NVIDIA-proprietary are not a contribution and subject to the terms and conditions at the top of the file
def tensor2vid(video: torch.Tensor, processor, output_type="np"):
# Based on:
# https://github.com/modelscope/modelscope/blob/1509fdb973e5871f37148a4b5e5964cafd43e64d/modelscope/pipelines/multi_modal/text_to_video_synthesis_pipeline.py#L78
batch_size, channels, num_frames, height, width = video.shape
outputs = []
for batch_idx in range(batch_size):
batch_vid = video[batch_idx].permute(1, 0, 2, 3)
batch_output = processor.postprocess(batch_vid, output_type)
outputs.append(batch_output)
return outputs
def prepare_mask_and_masked_image(image, mask):
"""
image: PIL.Image.Image
mask: PIL.Image.Image
"""
if isinstance(image, Image.Image):
image = np.array(image.convert("RGB"))
image = image[None].transpose(0, 3, 1, 2)
image = torch.from_numpy(image).to(dtype=torch.float32).contiguous() / 127.5 - 1.0
if isinstance(mask, Image.Image):
mask = np.array(mask.convert("L"))
mask = mask.astype(np.float32) / 255.0
mask = mask[None, None]
mask[mask < 0.5] = 0
mask[mask >= 0.5] = 1
mask = torch.from_numpy(mask).to(dtype=torch.float32).contiguous()
masked_image = image * (mask < 0.5)
return mask, masked_image
def download_image(url):
response = requests.get(url)
return Image.open(BytesIO(response.content)).convert("RGB")
def get_refit_weights(state_dict, onnx_opt_path, weight_name_mapping, weight_shape_mapping):
onnx_opt_dir = os.path.dirname(onnx_opt_path)
onnx_opt_model = onnx.load(onnx_opt_path)
# Create initializer data hashes
initializer_hash_mapping = {}
for initializer in onnx_opt_model.graph.initializer:
initializer_data = numpy_helper.to_array(initializer, base_dir=onnx_opt_dir).astype(np.float16)
initializer_hash = hash(initializer_data.data.tobytes())
initializer_hash_mapping[initializer.name] = initializer_hash
refit_weights = OrderedDict()
updated_weight_names = set() # save names of updated weights to refit only the required weights
for wt_name, wt in state_dict.items():
# query initializer to compare
initializer_name = weight_name_mapping[wt_name]
initializer_hash = initializer_hash_mapping[initializer_name]
# get shape transform info
initializer_shape, is_transpose = weight_shape_mapping[wt_name]
if is_transpose:
wt = torch.transpose(wt, 0, 1)
else:
wt = torch.reshape(wt, initializer_shape)
# include weight if hashes differ
wt_hash = hash(wt.cpu().detach().numpy().astype(np.float16).data.tobytes())
if initializer_hash != wt_hash:
updated_weight_names.add(initializer_name)
# Store all weights as the refitter may require unchanged weights too
# docs: https://docs.nvidia.com/deeplearning/tensorrt/developer-guide/index.html#refitting-engine-c
refit_weights[initializer_name] = wt.contiguous()
return refit_weights, updated_weight_names
def load_calib_prompts(batch_size, calib_data_path):
with open(calib_data_path, "r", encoding="utf-8") as file:
lst = [line.rstrip("\n") for line in file]
return [lst[i : i + batch_size] for i in range(0, len(lst), batch_size)]
def load_calibration_images(folder_path):
images = []
for filename in os.listdir(folder_path):
img_path = os.path.join(folder_path, filename)
if os.path.isfile(img_path):
image = load_image(img_path)
if image is not None:
images.append(image)
return images
class PercentileAmaxes:
def __init__(self, total_step, percentile) -> None:
self.data = {}
self.total_step = total_step
self.percentile = percentile
self.i = 0
def append(self, item):
_cur_step = self.i % self.total_step
if _cur_step not in self.data.keys():
self.data[_cur_step] = item
else:
self.data[_cur_step] = np.maximum(self.data[_cur_step], item)
self.i += 1
def add_arguments(parser):
# Stable Diffusion configuration
parser.add_argument('--version', type=str, default="1.5", choices=("1.4", "1.5", "dreamshaper-7", "2.0-base", "2.0", "2.1-base", "2.1", "xl-1.0", "xl-turbo", "svd-xt-1.1", "sd3", "cascade", "flux.1-dev", "flux.1-schnell", "flux.1-dev-canny", "flux.1-dev-depth"), help="Version of Stable Diffusion")
parser.add_argument('prompt', nargs = '*', help="Text prompt(s) to guide image generation")
parser.add_argument('--negative-prompt', nargs = '*', default=[''], help="The negative prompt(s) to guide the image generation.")
parser.add_argument('--batch-size', type=int, default=1, choices=[1, 2, 4], help="Batch size (repeat prompt)")
parser.add_argument('--batch-count', type=int, default=1, help="Number of images to generate in sequence, one at a time.")
parser.add_argument('--height', type=int, default=512, help="Height of image to generate (must be multiple of 8)")
parser.add_argument('--width', type=int, default=512, help="Height of image to generate (must be multiple of 8)")
parser.add_argument('--denoising-steps', type=int, default=30, help="Number of denoising steps")
parser.add_argument('--scheduler', type=str, default=None, choices=("DDIM", "DDPM", "EulerA", "Euler", "LCM", "LMSD", "PNDM", "UniPC", "DDPMWuerstchen", "FlowMatchEuler"), help="Scheduler for diffusion process")
parser.add_argument('--guidance-scale', type=float, default=7.5, help="Value of classifier-free guidance scale (must be greater than 1)")
parser.add_argument('--lora-scale', type=float, default=1.0, help="Controls how much to influence the outputs with the LoRA parameters. (must between 0 and 1)")
parser.add_argument('--lora-weight', type=float, nargs='+', default=None, help="The LoRA adapter(s) weights to use with the UNet. (must between 0 and 1)")
parser.add_argument('--lora-path', type=str, nargs='+', default=None, help="Path to LoRA adaptor. Ex: 'latent-consistency/lcm-lora-sdv1-5'")
# ONNX export
parser.add_argument('--onnx-opset', type=int, default=19, choices=range(7,20), help="Select ONNX opset version to target for exported models")
parser.add_argument('--onnx-dir', default='onnx', help="Output directory for ONNX export")
# Framework model ckpt
parser.add_argument('--framework-model-dir', default='pytorch_model', help="Directory for HF saved models")
# TensorRT engine build
parser.add_argument('--engine-dir', default='engine', help="Output directory for TensorRT engines")
parser.add_argument('--int8', action='store_true', help="Apply int8 quantization.")
parser.add_argument('--fp8', action='store_true', help="Apply fp8 quantization.")
parser.add_argument('--fp4', action='store_true', help="Apply fp4 quantization.")
parser.add_argument('--quantization-level', type=float, default=0.0, choices=[0.0, 1.0, 2.0, 2.5, 3.0, 4.0], help="int8/fp8 quantization level, 1: CNN, 2: CNN + FFN, 2.5: CNN + FFN + QKV, 3: CNN + Almost all Linear (Including FFN, QKV, Proj and others), 4: CNN + Almost all Linear + fMHA, 0: Default to 2.5 for int8 and 4.0 for fp8.")
parser.add_argument('--optimization-level', type=int, default=None, help=f"Set the builder optimization level to build the engine with. A higher level allows TensorRT to spend more building time for more optimization options. Must be one of {VALID_OPTIMIZATION_LEVELS}.")
parser.add_argument('--build-static-batch', action='store_true', help="Build TensorRT engines with fixed batch size.")
parser.add_argument('--build-dynamic-shape', action='store_true', help="Build TensorRT engines with dynamic image shapes.")
parser.add_argument('--build-enable-refit', action='store_true', help="Enable Refit option in TensorRT engines during build.")
parser.add_argument('--build-all-tactics', action='store_true', help="Build TensorRT engines using all tactic sources.")
parser.add_argument('--timing-cache', default=None, type=str, help="Path to the precached timing measurements to accelerate build.")
# TensorRT inference
parser.add_argument('--num-warmup-runs', type=int, default=5, help="Number of warmup runs before benchmarking performance")
parser.add_argument('--use-cuda-graph', action='store_true', help="Enable cuda graph")
parser.add_argument('--nvtx-profile', action='store_true', help="Enable NVTX markers for performance profiling")
parser.add_argument('--torch-inference', default='', help="Run inference with PyTorch (using specified compilation mode) instead of TensorRT.")
parser.add_argument('--seed', type=int, default=None, help="Seed for random generator to get consistent results")
parser.add_argument('--output-dir', default='output', help="Output directory for logs and image artifacts")
parser.add_argument('--hf-token', type=str, help="HuggingFace API access token for downloading model checkpoints")
parser.add_argument('-v', '--verbose', action='store_true', help="Show verbose output")
return parser
def process_pipeline_args(args):
if args.height % 8 != 0 or args.width % 8 != 0:
raise ValueError(f"Image height and width have to be divisible by 8 but specified as: {args.image_height} and {args.width}.")
max_batch_size = 4
if args.batch_size > max_batch_size:
raise ValueError(f"Batch size {args.batch_size} is larger than allowed {max_batch_size}.")
if args.use_cuda_graph and (not args.build_static_batch or args.build_dynamic_shape):
raise ValueError("Using CUDA graph requires static dimensions. Enable `--build-static-batch` and do not specify `--build-dynamic-shape`")
if args.optimization_level is None:
if args.int8 or args.fp8:
args.optimization_level = 4
else:
args.optimization_level = 3
if args.optimization_level not in VALID_OPTIMIZATION_LEVELS:
raise ValueError(f"Optimization level {args.optimization_level} not valid. Valid values are: {VALID_OPTIMIZATION_LEVELS}")
if args.int8 and not any(args.version.startswith(prefix) for prefix in ['xl', '1.4', '1.5', '2.1']):
raise ValueError("int8 quantization is only supported for SDXL, SD1.4, SD1.5 and SD2.1 pipelines.")
if args.fp8 and not any(args.version.startswith(prefix) for prefix in ('xl', '1.4', '1.5', '2.1', 'flux.1-dev', 'flux.1-schnell')):
raise ValueError("fp8 quantization is only supported for SDXL, SD1.4, SD1.5, SD2.1 and FLUX pipelines.")
if args.fp8 and args.int8:
raise ValueError("Cannot apply both int8 and fp8 quantization, please choose only one.")
if args.fp8:
device_info = torch.cuda.get_device_properties(0)
version = device_info.major * 10 + device_info.minor
if version < 89:
raise ValueError(f"Cannot apply FP8 quantization for GPU with compute capability {version / 10.0}. Only Ada and Hopper are supported.")
if args.quantization_level == 0.0:
def override_quant_level(level : float, dtype_str : str):
args.quantization_level = level
print(f"The default quantization level has been set to {level} for {dtype_str}.")
if args.fp8:
override_quant_level(3.0 if args.version in ("1.4", "1.5") else 4.0, "FP8")
elif args.int8:
override_quant_level(3.0, "INT8")
if args.fp4:
# FP4 precision is only supported for Flux Pipelines
assert args.version.startswith("flux"), "FP4 precision is only supported for Flux pipelines"
if args.lora_path and not any(args.version.startswith(prefix) for prefix in ('1.5', '2.1', 'xl')):
raise ValueError("LoRA adapter support is only supported for SD1.5, SD2.1 and SDXL pipelines")
if args.lora_weight:
for weight in (weight for weight in args.lora_weight if not 0 <= weight <= 1):
raise ValueError(f"LoRA adapter weights must be between 0 and 1, provided {weight}")
if not 0 <= args.lora_scale <= 1:
raise ValueError(f"LoRA scale value must be between 0 and 1, provided {args.lora_scale}")
kwargs_init_pipeline = {
'version': args.version,
'max_batch_size': max_batch_size,
'denoising_steps': args.denoising_steps,
'scheduler': args.scheduler,
'guidance_scale': args.guidance_scale,
'output_dir': args.output_dir,
'hf_token': args.hf_token,
'verbose': args.verbose,
'nvtx_profile': args.nvtx_profile,
'use_cuda_graph': args.use_cuda_graph,
'lora_scale': args.lora_scale,
'lora_weight': args.lora_weight,
'lora_path': args.lora_path,
'framework_model_dir': args.framework_model_dir,
'torch_inference': args.torch_inference,
}
kwargs_load_engine = {
'onnx_opset': args.onnx_opset,
'opt_batch_size': args.batch_size,
'opt_image_height': args.height,
'opt_image_width': args.width,
'optimization_level': args.optimization_level,
'static_batch': args.build_static_batch,
'static_shape': not args.build_dynamic_shape,
'enable_all_tactics': args.build_all_tactics,
'enable_refit': args.build_enable_refit,
'timing_cache': args.timing_cache,
'int8': args.int8,
'fp8': args.fp8,
'quantization_level': args.quantization_level,
}
args_run_demo = (args.prompt, args.negative_prompt, args.height, args.width, args.batch_size, args.batch_count, args.num_warmup_runs, args.use_cuda_graph)
return kwargs_init_pipeline, kwargs_load_engine, args_run_demo
docker/rockylinux8.Dockerfile
+12 -12
View File
@@ -25,7 +25,7 @@ ENV NV_CUDNN_VERSION 8.9.6.50-1
ENV NV_CUDNN_PACKAGE libcudnn8-${NV_CUDNN_VERSION}.cuda12.2
ENV NV_CUDNN_PACKAGE_DEV libcudnn8-devel-${NV_CUDNN_VERSION}.cuda12.2
ENV TRT_VERSION 10.8.0.43
ENV TRT_VERSION 10.9.0.34
SHELL ["/bin/bash", "-c"]
RUN dnf install -y \
@@ -62,19 +62,19 @@ RUN dnf install -y python38 python38-devel &&\
# Install TensorRT
RUN if [ "${CUDA_VERSION:0:2}" = "11" ]; then \
wget https://developer.nvidia.com/downloads/compute/machine-learning/tensorrt/10.8.0/tars/TensorRT-10.8.0.43.Linux.x86_64-gnu.cuda-11.8.tar.gz \
&& tar -xf TensorRT-10.8.0.43.Linux.x86_64-gnu.cuda-11.8.tar.gz \
&& cp -a TensorRT-10.8.0.43/lib/*.so* /usr/lib64 \
&& pip install TensorRT-10.8.0.43/python/tensorrt-10.8.0.43-cp38-none-linux_x86_64.whl ;\
elif [ "${CUDA_VERSION:0:2}" = "12" ]; then \
wget https://developer.nvidia.com/downloads/compute/machine-learning/tensorrt/10.8.0/tars/TensorRT-10.8.0.43.Linux.x86_64-gnu.cuda-12.8.tar.gz \
&& tar -xf TensorRT-10.8.0.43.Linux.x86_64-gnu.cuda-12.8.tar.gz \
&& cp -a TensorRT-10.8.0.43/lib/*.so* /usr/lib64 \
&& pip install TensorRT-10.8.0.43/python/tensorrt-10.8.0.43-cp38-none-linux_x86_64.whl ;\
else \
wget https://developer.nvidia.com/downloads/compute/machine-learning/tensorrt/10.9.0/tars/TensorRT-10.9.0.34.Linux.x86_64-gnu.cuda-11.8.tar.gz \
&& tar -xf TensorRT-10.9.0.34.Linux.x86_64-gnu.cuda-11.8.tar.gz \
&& cp -a TensorRT-10.9.0.34/lib/*.so* /usr/lib64 \
&& pip install TensorRT-10.9.0.34/python/tensorrt-10.9.0.34-cp38-none-linux_x86_64.whl ;\
elif [ "${CUDA_VERSION:0:2}" = "12" ]; then \
wget https://developer.nvidia.com/downloads/compute/machine-learning/tensorrt/10.9.0/tars/TensorRT-10.9.0.34.Linux.x86_64-gnu.cuda-12.8.tar.gz \
&& tar -xf TensorRT-10.9.0.34.Linux.x86_64-gnu.cuda-12.8.tar.gz \
&& cp -a TensorRT-10.9.0.34/lib/*.so* /usr/lib64 \
&& pip install TensorRT-10.9.0.34/python/tensorrt-10.9.0.34-cp38-none-linux_x86_64.whl ;\
else \
echo "Invalid CUDA_VERSION"; \
exit 1; \
fi
fi
# Install PyPI packages
RUN pip install --upgrade pip
docker/rockylinux9.Dockerfile
+12 -12
View File
@@ -25,7 +25,7 @@ ENV NV_CUDNN_VERSION 8.9.6.50-1
ENV NV_CUDNN_PACKAGE libcudnn8-${NV_CUDNN_VERSION}.cuda12.2
ENV NV_CUDNN_PACKAGE_DEV libcudnn8-devel-${NV_CUDNN_VERSION}.cuda12.2
ENV TRT_VERSION 10.8.0.43
ENV TRT_VERSION 10.9.0.34
SHELL ["/bin/bash", "-c"]
RUN dnf install -y \
@@ -67,19 +67,19 @@ RUN dnf -y install \
# Install TensorRT
RUN if [ "${CUDA_VERSION:0:2}" = "11" ]; then \
wget https://developer.nvidia.com/downloads/compute/machine-learning/tensorrt/10.8.0/tars/TensorRT-10.8.0.43.Linux.x86_64-gnu.cuda-11.8.tar.gz \
&& tar -xf TensorRT-10.8.0.43.Linux.x86_64-gnu.cuda-11.8.tar.gz \
&& cp -a TensorRT-10.8.0.43/lib/*.so* /usr/lib64 \
&& pip install TensorRT-10.8.0.43/python/tensorrt-10.8.0.43-cp39-none-linux_x86_64.whl ;\
elif [ "${CUDA_VERSION:0:2}" = "12" ]; then \
wget https://developer.nvidia.com/downloads/compute/machine-learning/tensorrt/10.8.0/tars/TensorRT-10.8.0.43.Linux.x86_64-gnu.cuda-12.8.tar.gz \
&& tar -xf TensorRT-10.8.0.43.Linux.x86_64-gnu.cuda-12.8.tar.gz \
&& cp -a TensorRT-10.8.0.43/lib/*.so* /usr/lib64 \
&& pip install TensorRT-10.8.0.43/python/tensorrt-10.8.0.43-cp39-none-linux_x86_64.whl ;\
else \
wget https://developer.nvidia.com/downloads/compute/machine-learning/tensorrt/10.9.0/tars/TensorRT-10.9.0.34.Linux.x86_64-gnu.cuda-11.8.tar.gz \
&& tar -xf TensorRT-10.9.0.34.Linux.x86_64-gnu.cuda-11.8.tar.gz \
&& cp -a TensorRT-10.9.0.34/lib/*.so* /usr/lib64 \
&& pip install TensorRT-10.9.0.34/python/tensorrt-10.9.0.34-cp39-none-linux_x86_64.whl ;\
elif [ "${CUDA_VERSION:0:2}" = "12" ]; then \
wget https://developer.nvidia.com/downloads/compute/machine-learning/tensorrt/10.9.0/tars/TensorRT-10.9.0.34.Linux.x86_64-gnu.cuda-12.8.tar.gz \
&& tar -xf TensorRT-10.9.0.34.Linux.x86_64-gnu.cuda-12.8.tar.gz \
&& cp -a TensorRT-10.9.0.34/lib/*.so* /usr/lib64 \
&& pip install TensorRT-10.9.0.34/python/tensorrt-10.9.0.34-cp39-none-linux_x86_64.whl ;\
else \
echo "Invalid CUDA_VERSION"; \
exit 1; \
fi
fi
# Install Cmake
RUN cd /tmp && \
docker/ubuntu-20.04.Dockerfile
+16 -16
View File
@@ -28,7 +28,7 @@ ENV CUDA_VERSION_MAJOR_MINOR=12.2
ENV NV_CUDNN_PACKAGE "libcudnn8=$NV_CUDNN_VERSION-1+cuda${CUDA_VERSION_MAJOR_MINOR}"
ENV NV_CUDNN_PACKAGE_DEV "libcudnn8-dev=$NV_CUDNN_VERSION-1+cuda${CUDA_VERSION_MAJOR_MINOR}"
ENV TRT_VERSION 10.8.0.43
ENV TRT_VERSION 10.9.0.34
SHELL ["/bin/bash", "-c"]
RUN apt-get update && apt-get install -y --no-install-recommends \
@@ -74,29 +74,29 @@ RUN apt-get update && apt-get install -y --no-install-recommends \
# Install python3
RUN apt-get install -y --no-install-recommends \
python3 \
python3-pip \
python3-dev \
python3-wheel &&\
python3 \
python3-pip \
python3-dev \
python3-wheel &&\
cd /usr/local/bin &&\
ln -s /usr/bin/python3 python &&\
ln -s /usr/bin/pip3 pip;
# Install TensorRT
RUN if [ "${CUDA_VERSION:0:2}" = "11" ]; then \
wget https://developer.nvidia.com/downloads/compute/machine-learning/tensorrt/10.8.0/tars/TensorRT-10.8.0.43.Linux.x86_64-gnu.cuda-11.8.tar.gz \
&& tar -xf TensorRT-10.8.0.43.Linux.x86_64-gnu.cuda-11.8.tar.gz \
&& cp -a TensorRT-10.8.0.43/lib/*.so* /usr/lib/x86_64-linux-gnu \
&& pip install TensorRT-10.8.0.43/python/tensorrt-10.8.0.43-cp38-none-linux_x86_64.whl ;\
elif [ "${CUDA_VERSION:0:2}" = "12" ]; then \
wget https://developer.nvidia.com/downloads/compute/machine-learning/tensorrt/10.8.0/tars/TensorRT-10.8.0.43.Linux.x86_64-gnu.cuda-12.8.tar.gz \
&& tar -xf TensorRT-10.8.0.43.Linux.x86_64-gnu.cuda-12.8.tar.gz \
&& cp -a TensorRT-10.8.0.43/lib/*.so* /usr/lib/x86_64-linux-gnu \
&& pip install TensorRT-10.8.0.43/python/tensorrt-10.8.0.43-cp38-none-linux_x86_64.whl ;\
else \
wget https://developer.nvidia.com/downloads/compute/machine-learning/tensorrt/10.9.0/tars/TensorRT-10.9.0.34.Linux.x86_64-gnu.cuda-11.8.tar.gz \
&& tar -xf TensorRT-10.9.0.34.Linux.x86_64-gnu.cuda-11.8.tar.gz \
&& cp -a TensorRT-10.9.0.34/lib/*.so* /usr/lib/x86_64-linux-gnu \
&& pip install TensorRT-10.9.0.34/python/tensorrt-10.9.0.34-cp38-none-linux_x86_64.whl ;\
elif [ "${CUDA_VERSION:0:2}" = "12" ]; then \
wget https://developer.nvidia.com/downloads/compute/machine-learning/tensorrt/10.9.0/tars/TensorRT-10.9.0.34.Linux.x86_64-gnu.cuda-12.8.tar.gz \
&& tar -xf TensorRT-10.9.0.34.Linux.x86_64-gnu.cuda-12.8.tar.gz \
&& cp -a TensorRT-10.9.0.34/lib/*.so* /usr/lib/x86_64-linux-gnu \
&& pip install TensorRT-10.9.0.34/python/tensorrt-10.9.0.34-cp38-none-linux_x86_64.whl ;\
else \
echo "Invalid CUDA_VERSION"; \
exit 1; \
fi
fi
# Install PyPI packages
RUN pip3 install --upgrade pip
docker/ubuntu-22.04-aarch64.Dockerfile
+10 -10
View File
@@ -20,7 +20,7 @@ ARG CUDA_VERSION=12.8.0
# Multi-arch container support available in non-cudnn containers.
FROM nvidia/cuda:${CUDA_VERSION}-devel-ubuntu22.04
ENV TRT_VERSION 10.8.0.43
ENV TRT_VERSION 10.9.0.34
SHELL ["/bin/bash", "-c"]
# Setup user account
@@ -60,10 +60,10 @@ RUN apt-get update && apt-get install -y --no-install-recommends \
# Install python3
RUN apt-get install -y --no-install-recommends \
python3 \
python3-pip \
python3-dev \
python3-wheel &&\
python3 \
python3-pip \
python3-dev \
python3-wheel &&\
cd /usr/local/bin &&\
ln -s /usr/bin/python3 python &&\
ln -s /usr/bin/pip3 pip;
@@ -71,16 +71,16 @@ RUN apt-get install -y --no-install-recommends \
# Install TensorRT. This will also pull in CUDNN
RUN ver="${CUDA_VERSION%.*}" &&\
if [ "${ver%.*}" = "12" ] ; then \
ver="12.8"; \
ver="12.8"; \
fi &&\
v="${TRT_VERSION}-1+cuda${ver}" &&\
apt-key adv --fetch-keys https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2204/sbsa/3bf863cc.pub &&\
apt-get update &&\
sudo apt-get -y install libnvinfer10=${v} libnvonnxparsers10=${v} libnvinfer-plugin10=${v} \
libnvinfer-dev=${v} libnvonnxparsers-dev=${v} libnvinfer-plugin-dev=${v} \
python3-libnvinfer=${v} libnvinfer-dispatch10=${v} libnvinfer-dispatch-dev=${v} libnvinfer-lean10=${v} \
libnvinfer-lean-dev=${v} libnvinfer-vc-plugin10=${v} libnvinfer-vc-plugin-dev=${v} \
libnvinfer-headers-dev=${v} libnvinfer-headers-plugin-dev=${v};
libnvinfer-dev=${v} libnvonnxparsers-dev=${v} libnvinfer-plugin-dev=${v} \
python3-libnvinfer=${v} libnvinfer-dispatch10=${v} libnvinfer-dispatch-dev=${v} libnvinfer-lean10=${v} \
libnvinfer-lean-dev=${v} libnvinfer-vc-plugin10=${v} libnvinfer-vc-plugin-dev=${v} \
libnvinfer-headers-dev=${v} libnvinfer-headers-plugin-dev=${v};
# Install Cmake
RUN cd /tmp && \
docker/ubuntu-22.04.Dockerfile
+16 -16
View File
@@ -28,7 +28,7 @@ ENV CUDA_VERSION_MAJOR_MINOR=12.2
ENV NV_CUDNN_PACKAGE "libcudnn8=$NV_CUDNN_VERSION-1+cuda${CUDA_VERSION_MAJOR_MINOR}"
ENV NV_CUDNN_PACKAGE_DEV "libcudnn8-dev=$NV_CUDNN_VERSION-1+cuda${CUDA_VERSION_MAJOR_MINOR}"
ENV TRT_VERSION 10.8.0.43
ENV TRT_VERSION 10.9.0.34
SHELL ["/bin/bash", "-c"]
RUN apt-get update && apt-get install -y --no-install-recommends \
@@ -74,29 +74,29 @@ RUN apt-get update && apt-get install -y --no-install-recommends \
# Install python3
RUN apt-get install -y --no-install-recommends \
python3 \
python3-pip \
python3-dev \
python3-wheel &&\
python3 \
python3-pip \
python3-dev \
python3-wheel &&\
cd /usr/local/bin &&\
ln -s /usr/bin/python3 python &&\
ln -s /usr/bin/pip3 pip;
# Install TensorRT
RUN if [ "${CUDA_VERSION:0:2}" = "11" ]; then \
wget https://developer.nvidia.com/downloads/compute/machine-learning/tensorrt/10.8.0/tars/TensorRT-10.8.0.43.Linux.x86_64-gnu.cuda-11.8.tar.gz \
&& tar -xf TensorRT-10.8.0.43.Linux.x86_64-gnu.cuda-11.8.tar.gz \
&& cp -a TensorRT-10.8.0.43/lib/*.so* /usr/lib/x86_64-linux-gnu \
&& pip install TensorRT-10.8.0.43/python/tensorrt-10.8.0.43-cp310-none-linux_x86_64.whl ;\
elif [ "${CUDA_VERSION:0:2}" = "12" ]; then \
wget https://developer.nvidia.com/downloads/compute/machine-learning/tensorrt/10.8.0/tars/TensorRT-10.8.0.43.Linux.x86_64-gnu.cuda-12.8.tar.gz \
&& tar -xf TensorRT-10.8.0.43.Linux.x86_64-gnu.cuda-12.8.tar.gz \
&& cp -a TensorRT-10.8.0.43/lib/*.so* /usr/lib/x86_64-linux-gnu \
&& pip install TensorRT-10.8.0.43/python/tensorrt-10.8.0.43-cp310-none-linux_x86_64.whl ;\
else \
wget https://developer.nvidia.com/downloads/compute/machine-learning/tensorrt/10.9.0/tars/TensorRT-10.9.0.34.Linux.x86_64-gnu.cuda-11.8.tar.gz \
&& tar -xf TensorRT-10.9.0.34.Linux.x86_64-gnu.cuda-11.8.tar.gz \
&& cp -a TensorRT-10.9.0.34/lib/*.so* /usr/lib/x86_64-linux-gnu \
&& pip install TensorRT-10.9.0.34/python/tensorrt-10.9.0.34-cp310-none-linux_x86_64.whl ;\
elif [ "${CUDA_VERSION:0:2}" = "12" ]; then \
wget https://developer.nvidia.com/downloads/compute/machine-learning/tensorrt/10.9.0/tars/TensorRT-10.9.0.34.Linux.x86_64-gnu.cuda-12.8.tar.gz \
&& tar -xf TensorRT-10.9.0.34.Linux.x86_64-gnu.cuda-12.8.tar.gz \
&& cp -a TensorRT-10.9.0.34/lib/*.so* /usr/lib/x86_64-linux-gnu \
&& pip install TensorRT-10.9.0.34/python/tensorrt-10.9.0.34-cp310-none-linux_x86_64.whl ;\
else \
echo "Invalid CUDA_VERSION"; \
exit 1; \
fi
fi
# Install PyPI packages
RUN pip3 install --upgrade pip
docker/ubuntu-cross-aarch64.Dockerfile
+4 -4
View File
@@ -21,7 +21,7 @@ ARG OS_VERSION=22.04
FROM nvidia/cuda:${CUDA_VERSION}-devel-ubuntu${OS_VERSION}
LABEL maintainer="NVIDIA CORPORATION"
ENV TRT_VERSION 10.8.0.43
ENV TRT_VERSION 10.9.0.34
ENV DEBIAN_FRONTEND=noninteractive
ARG uid=1000
@@ -120,9 +120,9 @@ RUN cd /pdk_files/tensorrt \
&& cd lib \
&& mkdir stubs \
&& for x in nvinfer nvparsers nvinfer_plugin nvonnxparser; \
do \
CC=aarch64-linux-gnu-gcc /pdk_files/stubify.sh lib${x}.so stubs/lib${x}.so \
; done
do \
CC=aarch64-linux-gnu-gcc /pdk_files/stubify.sh lib${x}.so stubs/lib${x}.so \
; done
# Set environment and working directory
ENV TRT_LIBPATH /pdk_files/tensorrt/lib
include/NvInfer.h
+52 -26
View File
@@ -29,7 +29,7 @@
//!
//! Please see the accompanying user guide and samples for higher-level information and general advice on
//! using TensorRT.
//
//!
//! TensorRT Versioning follows Semantic Versioning Guidelines specified here: https://semver.org/
//!
@@ -106,7 +106,7 @@ enum class LayerType : int32_t
kSQUEEZE = 47, //!< Squeeze Layer.
kUNSQUEEZE = 48, //!< Unsqueeze Layer.
kCUMULATIVE = 49, //!< Cumulative layer.
kDYNAMIC_QUANTIZE = 50 //!< Dynamic Quantize layer.
kDYNAMIC_QUANTIZE = 50, //!< Dynamic Quantize layer.
};
//!
@@ -3416,7 +3416,7 @@ public:
//! \param index the index of the input to modify.
//! \param tensor the new input tensor
//!
//! For a slice layer, the values 0-4 are valid.
//! For a slice layer, the values 0-5 are valid.
//! The indices are as follows:
//!
//! - 0: Tensor to be sliced.
@@ -4357,7 +4357,9 @@ protected:
};
//!
//! \enum Enum that describes kinds of loop outputs.
//! \enum LoopOutput
//!
//! \brief Enum that describes kinds of loop outputs.
//!
enum class LoopOutput : int32_t
{
@@ -4383,7 +4385,9 @@ constexpr inline int32_t EnumMax<LoopOutput>() noexcept
}
//!
//! \enum Enum that describes kinds of trip limits.
//! \enum TripLimit
//!
//! \brief Enum that describes kinds of trip limits.
//!
enum class TripLimit : int32_t
{
@@ -6407,6 +6411,7 @@ protected:
virtual ~INormalizationLayer() noexcept = default;
};
//!
//! \class ISqueezeLayer
//!
@@ -7745,7 +7750,7 @@ public:
//!
//! \see IDequantizeLayer
//!
//! \p input tensor data type must be DataType::kINT8/DataType::kFP8.
//! \p input tensor data type must be DataType::kINT8 or DataType::kFP8.
//! \p scale tensor data type must be DataType::kFLOAT. The subgraph which terminates with the \p scale tensor must
//! be a build-time constant.
//!
@@ -7763,10 +7768,11 @@ public:
//!
//! \param input The input tensor to be dequantized.
//! \param scale A tensor with the scale value.
//! \param outputType Output tensor data type.
//!
//! \see IDequantizeLayer
//!
//! \p input tensor data type must be DataType::kINT8/DataType::kFP8/DataType::kINT4.
//! \p input tensor data type must be DataType::kINT8, DataType::kFP8 or DataType::kINT4.
//! \p scale tensor data type defaults to DataType::kFLOAT. For strongly typed networks, it must be the same as the
//! output data type. The subgraph which terminates with the \p scale tensor must be a build-time constant.
//! \p outputType output tensor data type, default value is DataType::kFLOAT. Future calls to set output type using
@@ -7807,7 +7813,7 @@ public:
//!
//! \see IQuantizeLayer
//!
//! \p input tensor data type must be DataType::kFLOAT/DataType::kHALF.
//! \p input tensor data type must be DataType::kFLOAT or DataType::kHALF.
//! \p scale tensor data type must be DataType::kFLOAT. The subgraph which terminates with the \p scale tensor must
//! be a build-time constant.
//!
@@ -7825,10 +7831,11 @@ public:
//!
//! \param input The input tensor to be quantized.
//! \param scale A tensor with the scale value.
//! \param outputType Output tensor data type.
//!
//! \see IQuantizeLayer
//!
//! \p input tensor data type must be DataType::kFLOAT/DataType::kHALF/DataType::kBF16.
//! \p input tensor data type must be DataType::kFLOAT, DataType::kHALF or DataType::kBF16.
//! \p scale tensor data type defaults to DataType::kFLOAT. For strongly typed networks, it must have the same data
//! type as the input. The subgraph which terminates with the \p scale tensor must be a build-time constant.
//! \p outputType output tensor data type, must be DataType::kINT8 (default), DataType::kFP8 or DataType::kINT4.
@@ -7853,10 +7860,9 @@ public:
//! \param axis The axis that is sliced into blocks. The axis must be the last or second to last dimension.
//! \param blockSize The number of elements that are quantized using a shared scale factor.
//! Currently only blocks of 16 elements are supported.
//!
//! \p outputType The data type of the quantized output tensor, must be DataType::kFP4. Future calls to set output
//! \param outputType The data type of the quantized output tensor, must be DataType::kFP4. Future calls to set output
//! type using setToType or setOutputType must be consistent.
//! \p scaleType The data type of the scale factor used for quantizing the input data, must be DataType::kFP8.
//! \param scaleType The data type of the scale factor used for quantizing the input data, must be DataType::kFP8.
//!
//! \return The new dynamic quantization layer, or nullptr if it could not be created.
//!
@@ -8813,7 +8819,7 @@ enum class BuilderFlag : int32_t
//! Enable plugins with FP8 input/output.
//!
//! This flag is not supported with hardware-compatibility mode.
//! This flag is not supported when HardwareCompatibilityLevel::kAMPERE_PLUS is enabled.
//!
//! \see HardwareCompatibilityLevel
kFP8 = 15,
@@ -9192,10 +9198,8 @@ struct EnumMaxImpl<PreviewFeature>
//!
//! \enum HardwareCompatibilityLevel
//!
//! \brief Describes requirements of compatibility with GPU architectures other than that of the GPU on which the engine was
//! built.
//!
//! Levels except kNONE are only supported for engines built on NVIDIA Ampere and later GPUs.
//! \brief Describes requirements of compatibility with GPU architectures other than that of the GPU on which the engine
//! was built.
//!
//! \warning Note that compatibility with future hardware depends on CUDA forward compatibility support.
//!
@@ -9209,12 +9213,24 @@ enum class HardwareCompatibilityLevel : int32_t
//! reserved and backend kernel max shared memory to 48KiB, may reduce the number of available tactics for each
//! layer, and may prevent some fusions from occurring. Thus this can decrease the performance, especially for tf32
//! models.
//! This option will disable cuDNN, cuBLAS, and cuBLAS LT as tactic sources.
//! This option will disable cuDNN, cuBLAS, and cuBLASLt as tactic sources.
//!
//! This option is only supported for engines built on NVIDIA Ampere and later GPUs.
//!
//! The driver reserved shared memory can be queried from cuDeviceGetAttribute(&reservedShmem,
//! CU_DEVICE_ATTRIBUTE_RESERVED_SHARED_MEMORY_PER_BLOCK).
//!
kAMPERE_PLUS = 1,
//! Require that the engine is compatible with GPUs that have the same Compute Capability
//! (https://developer.nvidia.com/cuda-gpus) as the one it was built on. This may decrease the performance compared
//! to an engine with no compatibility.
//!
//! This option will disable cuDNN, cuBLAS, and cuBLASLt as tactic sources.
//!
//! This option is only supported for engines built on NVIDIA Turing and later GPUs.
//!
kSAME_COMPUTE_CAPABILITY = 2,
};
namespace impl
@@ -9227,7 +9243,7 @@ namespace impl
template <>
struct EnumMaxImpl<HardwareCompatibilityLevel>
{
static constexpr int32_t kVALUE = 2;
static constexpr int32_t kVALUE = 3;
};
} // namespace impl
@@ -9259,9 +9275,9 @@ enum class TilingOptimizationLevel : int32_t
namespace impl
{
//!
//! Maximum number of elements in HardwareCompatibilityLevel enum.
//! Maximum number of elements in TilingOptimizationLevel enum.
//!
//! \see HardwareCompatibilityLevel
//! \see TilingOptimizationLevel
//!
template <>
struct EnumMaxImpl<TilingOptimizationLevel>
@@ -10308,6 +10324,15 @@ enum class NetworkDefinitionCreationFlag : int32_t
//! inputs/operator annotations. Setting layer precision and layer output types is not allowed, and the network
//! output types will be inferred based on the input types and the type inference rules.
kSTRONGLY_TYPED = 1,
//! If set, for a Python plugin with both AOT and JIT implementations, the JIT implementation will be used.
//! Any plugin-specific JIT/AOT specification may override this.
//! Cannot be used in conjunction with NetworkDefinitionCreationFlag::kPREFER_AOT_PYTHON_PLUGINS.
kPREFER_JIT_PYTHON_PLUGINS = 2,
//! If set, for a Python plugin with both AOT and JIT implementations, the AOT implementation will be used.
//! Any plugin-specific JIT/AOT specification may override this.
//! Cannot be used in conjunction with NetworkDefinitionCreationFlag::kPREFER_JIT_PYTHON_PLUGINS.
kPREFER_AOT_PYTHON_PLUGINS = 3,
};
//!
@@ -10318,7 +10343,7 @@ enum class NetworkDefinitionCreationFlag : int32_t
template <>
constexpr inline int32_t EnumMax<NetworkDefinitionCreationFlag>() noexcept
{
return 2;
return 4;
}
//!
@@ -10379,7 +10404,8 @@ public:
//! \param allocator Set the GPU allocator to be used by the builder. All GPU memory acquired will use this
//! allocator. If NULL is passed, the default allocator will be used.
//!
//! Default: uses cudaMalloc/cudaFree.
//! Default: allocateAsync uses cudaMallocAsync if cudaDevAttrMemoryPoolsSupported returns true, otherwise falls
//! back to cudaMalloc. allocate always uses cudaMalloc.
//!
//! \note This allocator will be passed to any engines created via the builder; thus the lifetime of the allocator
//! must span the lifetime of those engines as
@@ -10395,7 +10421,7 @@ public:
//!
//! \see IBuilderConfig
//!
[[nodiscard]] nvinfer1::IBuilderConfig* createBuilderConfig() noexcept
nvinfer1::IBuilderConfig* createBuilderConfig() noexcept
{
return mImpl->createBuilderConfig();
}
@@ -10417,7 +10443,7 @@ public:
//!
//! \see INetworkDefinition, NetworkDefinitionCreationFlags
//!
[[nodiscard]] nvinfer1::INetworkDefinition* createNetworkV2(NetworkDefinitionCreationFlags flags) noexcept
nvinfer1::INetworkDefinition* createNetworkV2(NetworkDefinitionCreationFlags flags) noexcept
{
return mImpl->createNetworkV2(flags);
}
@@ -10432,7 +10458,7 @@ public:
//!
//! \see IOptimizationProfile
//!
[[nodiscard]] nvinfer1::IOptimizationProfile* createOptimizationProfile() noexcept
nvinfer1::IOptimizationProfile* createOptimizationProfile() noexcept
{
return mImpl->createOptimizationProfile();
}
include/NvInferImpl.h
+5 -5
View File
@@ -1,5 +1,5 @@
/*
* SPDX-FileCopyrightText: Copyright (c) 1993-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
* SPDX-FileCopyrightText: Copyright (c) 1993-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
* SPDX-License-Identifier: Apache-2.0
*
* Licensed under the Apache License, Version 2.0 (the "License");
@@ -1039,6 +1039,7 @@ public:
virtual bool getReverse() const noexcept = 0;
}; // class VCumulativeLayer
class VNetworkDefinition : public VRoot
{
public:
@@ -1114,8 +1115,6 @@ public:
virtual INormalizationLayer* addNormalization(
ITensor& input, ITensor& scale, ITensor& bias, uint32_t axesMask) noexcept = 0;
virtual ICastLayer* addCast(ITensor& input, DataType toType) noexcept = 0;
virtual ICumulativeLayer* addCumulative(
ITensor& input, ITensor& axis, CumulativeOperation operation, bool exclusive, bool reverse) noexcept = 0;
virtual IBuilder& getBuilder() const noexcept = 0;
virtual NetworkDefinitionCreationFlags getFlags() const noexcept = 0;
virtual bool getFlag(NetworkDefinitionCreationFlag networkDefinitionCreationFlag) const noexcept = 0;
@@ -1131,8 +1130,9 @@ public:
virtual ISqueezeLayer* addSqueeze(ITensor& input, ITensor& axes) noexcept = 0;
virtual IUnsqueezeLayer* addUnsqueeze(ITensor& input, ITensor& axes) noexcept = 0;
virtual IDynamicQuantizeLayer* addDynamicQuantize(
ITensor& input, int32_t axis, int32_t blockSize, DataType toType, DataType scaleType) noexcept
= 0;
ITensor& input, int32_t axis, int32_t blockSize, DataType toType, DataType scaleType) noexcept = 0;
virtual ICumulativeLayer* addCumulative(
ITensor& input, ITensor& axis, CumulativeOperation operation, bool exclusive, bool reverse) noexcept = 0;
};
class VAlgorithmIOInfo : public VRoot
include/NvInferRuntime.h
+10 -9
View File
@@ -1,5 +1,5 @@
/*
* SPDX-FileCopyrightText: Copyright (c) 1993-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
* SPDX-FileCopyrightText: Copyright (c) 1993-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
* SPDX-License-Identifier: Apache-2.0
*
* Licensed under the Apache License, Version 2.0 (the "License");
@@ -899,6 +899,7 @@ public:
//! \param shapeInputs Expressions for values of the shape tensor inputs
//! \param nbShapeInputs The number of shape tensor inputs
//! \param outputs Pre-allocated array to which the output dimensions must be written
//! \param nbOutputs Number of outputs.
//! \param exprBuilder Object for generating new dimension expressions
//!
//! \note Any size tensor outputs must be declared to be 0D.
@@ -1384,13 +1385,14 @@ enum class TensorFormat : int32_t
//! The stride of each dimension is the product of the dimensions after it.
//! The last dimension has unit stride.
//!
//! This format supports all TensorRT types.
//! For DLA usage, the tensor sizes are limited to C,H,W in the range [1,8192].
kLINEAR = 0,
//! Vector-major format with two scalars per vector.
//! Vector dimension is third to last.
//!
//! This format requires FP16 and at least three dimensions.
//! This format requires FP16 or BF16 and at least three dimensions.
kCHW2 = 1,
//! Vector-minor format with eight scalars per vector.
@@ -1401,7 +1403,7 @@ enum class TensorFormat : int32_t
//! Vector-major format with four scalars per vector.
//! Vector dimension is third to last.
//!
//! This format requires INT8 or FP16 and at least three dimensions.
//! This format requires INT8, FP16, or BF16 and at least three dimensions.
//! For INT8, the length of the vector dimension must be a build-time constant.
//!
//! Deprecated usage:
@@ -1425,7 +1427,7 @@ enum class TensorFormat : int32_t
//! Vector-major format with 32 scalars per vector.
//! Vector dimension is third to last.
//!
//! This format requires at least three dimensions.
//! This format requires INT8, FP32, or FP16 and at least three dimensions.
//!
//! For DLA usage, this format maps to the native feature format for INT8,
//! and the tensor sizes are limited to C,H,W in the range [1,8192].
@@ -1445,7 +1447,7 @@ enum class TensorFormat : int32_t
//! Vector-minor format where channel dimension is third to last and unpadded.
//!
//! This format requires either FP32, FP16, UINT8, INT64 or BF16 and at least three dimensions.
//! This format requires either FP32 or UINT8 and at least three dimensions.
kHWC = 8,
//! DLA planar format. For a tensor with dimension {N, C, H, W}, the W axis
@@ -1858,7 +1860,8 @@ public:
//! \param allocator Set the GPU allocator to be used by the runtime. All GPU memory acquired will use this
//! allocator. If NULL is passed, the default allocator will be used.
//!
//! Default: uses cudaMalloc/cudaFree.
//! Default: allocateAsync uses cudaMallocAsync if cudaDevAttrMemoryPoolsSupported returns true, otherwise falls
//! back to cudaMalloc. allocate always uses cudaMalloc.
//!
//! If nullptr is passed, the default allocator will be used.
//!
@@ -1948,7 +1951,6 @@ public:
//! weight streaming is enabled.
//!
//! \param streamReader a read-only stream from which TensorRT will deserialize a previously serialized engine.
//! \param stream The CUDA stream used when performing asynchronous I/O.
//!
//! \return The engine, or nullptr if it could not be deserialized. The pointer may not be valid immediately after
//! the function returns.
@@ -3596,8 +3598,6 @@ public:
//! \return hardwareCompatibilityLevel The level of hardware
//! compatibility.
//!
//! This is only supported for Ampere and newer architectures.
//!
HardwareCompatibilityLevel getHardwareCompatibilityLevel() const noexcept
{
return mImpl->getHardwareCompatibilityLevel();
@@ -5092,6 +5092,7 @@ inline IRuntime* createInferRuntime(ILogger& logger) noexcept
//!
//! \brief Create an instance of an IRefitter class.
//!
//! \param engine The engine class for the refitter.
//! \param logger The logging class for the refitter.
//!
inline IRefitter* createInferRefitter(ICudaEngine& engine, ILogger& logger) noexcept
include/NvInferRuntimeBase.h
+3 -2
View File
@@ -1,5 +1,5 @@
/*
* SPDX-FileCopyrightText: Copyright (c) 1993-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
* SPDX-FileCopyrightText: Copyright (c) 1993-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
* SPDX-License-Identifier: Apache-2.0
*
* Licensed under the Apache License, Version 2.0 (the "License");
@@ -176,6 +176,7 @@ enum class DataType : int32_t
//! 4-bit floating point type
//! 1 bit sign, 2 bit exponent, 1 bit mantissa
kFP4 = 10,
};
namespace impl
@@ -184,7 +185,7 @@ namespace impl
template <>
struct EnumMaxImpl<DataType>
{
//! Declaration of kVALUE that represents the maximum number of elements in the DataType enum.
//! Declaration of kVALUE that represents the maximum number of elements in the DataType enum.
static constexpr int32_t kVALUE = 11;
};
} // namespace impl
include/NvInferVersion.h
+3 -3
View File
@@ -1,5 +1,5 @@
/*
* SPDX-FileCopyrightText: Copyright (c) 1993-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
* SPDX-FileCopyrightText: Copyright (c) 1993-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
* SPDX-License-Identifier: Apache-2.0
*
* Licensed under the Apache License, Version 2.0 (the "License");
@@ -24,9 +24,9 @@
#define NV_INFER_VERSION_H
#define NV_TENSORRT_MAJOR 10 //!< TensorRT major version.
#define NV_TENSORRT_MINOR 8 //!< TensorRT minor version.
#define NV_TENSORRT_MINOR 9 //!< TensorRT minor version.
#define NV_TENSORRT_PATCH 0 //!< TensorRT patch version.
#define NV_TENSORRT_BUILD 43 //!< TensorRT build number.
#define NV_TENSORRT_BUILD 34 //!< TensorRT build number.
#define NV_TENSORRT_LWS_MAJOR 0 //!< TensorRT LWS major version.
#define NV_TENSORRT_LWS_MINOR 0 //!< TensorRT LWS minor version.
parsers/CMakeLists.txt
+11 -4
View File
@@ -1,5 +1,5 @@
#
# SPDX-FileCopyrightText: Copyright (c) 1993-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-FileCopyrightText: Copyright (c) 1993-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#
# Licensed under the Apache License, Version 2.0 (the "License");
@@ -19,10 +19,17 @@
add_custom_target(parsers DEPENDS nvonnxparser)
add_definitions("-D_PROTOBUF_INSTALL_DIR=${Protobuf_INSTALL_DIR}")
add_compile_options("-Dgoogle=google_private")
add_compile_options("-Dgoogle=google_trtrepack")
set(TENSORRT_ROOT ${PROJECT_SOURCE_DIR})
set(TENSORRT_BUILD ${TRT_OUT_DIR} ${TRT_LIB_DIR})
set(CMAKE_LIBRARY_OUTPUT_DIRECTORY ${TRT_OUT_DIR})
set(TENSORRT_BUILD ${TRT_OUT_DIR})
# Write the libs out to the top-level of the build directory.
set(CMAKE_LIBRARY_OUTPUT_DIRECTORY ${TENSORRT_BUILD})
message(NOTICE "CMAKE_LIBRARY_OUTPUT_DIRECTORY ${CMAKE_LIBRARY_OUTPUT_DIRECTORY}")
# Write the libs out to the top-level of the build directory.
set(CMAKE_ARCHIVE_OUTPUT_DIRECTORY ${TENSORRT_BUILD})
message(NOTICE "CMAKE_ARCHIVE_OUTPUT_DIRECTORY ${CMAKE_ARCHIVE_OUTPUT_DIRECTORY}")
include_directories(
${Protobuf_INCLUDE_DIR}
parsers/onnx
+1 -1
Submodule parsers/onnx updated: 118ed0aea1...d5dce67db7
plugin/CMakeLists.txt
+236 -7
View File
@@ -1,5 +1,5 @@
#
# SPDX-FileCopyrightText: Copyright (c) 1993-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-FileCopyrightText: Copyright (c) 1993-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#
# Licensed under the Apache License, Version 2.0 (the "License");
@@ -14,6 +14,206 @@
# See the License for the specific language governing permissions and
# limitations under the License.
#
if (${TRT_BUILD_ENABLE_NEW_PLUGIN_FLOW})
option(TRT_BUILD_INCLUDE_BERT_QKV_PLUGIN "Build the BERT QKV to Context Plugin and related plugins." ON)
set(TRT_PLUGIN_NAMES
batchedNMSPlugin
batchTilePlugin
clipPlugin
coordConvACPlugin
cropAndResizePlugin
decodeBbox3DPlugin
detectionLayerPlugin
disentangledAttentionPlugin
efficientNMSPlugin
flattenConcat
generateDetectionPlugin
gridAnchorPlugin
groupNormalizationPlugin
instanceNormalizationPlugin
leakyReluPlugin
modulatedDeformConvPlugin
multilevelCropAndResizePlugin
multilevelProposeROI
multiscaleDeformableAttnPlugin
nmsPlugin
normalizePlugin
nvFasterRCNN
pillarScatterPlugin
priorBoxPlugin
proposalLayerPlugin
proposalPlugin
pyramidROIAlignPlugin
regionPlugin
reorgPlugin
resizeNearestPlugin
roiAlignPlugin
scatterElementsPlugin
scatterPlugin
specialSlicePlugin
splitPlugin
voxelGeneratorPlugin
)
if(${TRT_BUILD_INCLUDE_BERT_QKV_PLUGIN})
list(APPEND TRT_PLUGIN_NAMES
bertQKVToContextPlugin
embLayerNormPlugin
fcPlugin
geluPlugin
skipLayerNormPlugin
)
endif()
add_library(tensorrt_plugins SHARED)
add_library(tensorrt_plugins_internal SHARED)
function(add_plugin_source)
target_sources(tensorrt_plugins PRIVATE ${ARGN})
target_sources(tensorrt_plugins_internal PRIVATE ${ARGN})
endfunction()
add_library(tensorrt_vc_plugins SHARED)
function(add_vc_plugin_source)
target_sources(tensorrt_vc_plugins PRIVATE ${ARGN})
endfunction()
add_subdirectory(api)
add_subdirectory(vc)
add_subdirectory(common)
foreach(PLUGIN_NAME IN LISTS TRT_PLUGIN_NAMES)
add_subdirectory(${PLUGIN_NAME})
endforeach()
if (NOT MSVC)
set(trt_plugins_link_options
"LINKER:--version-script=${CMAKE_CURRENT_LIST_DIR}/exports.map"
"LINKER:-z,relro"
"LINKER:-Bsymbolic"
"LINKER:--no-undefined"
"LINKER:--no-as-needed"
"$<$<CONFIG:Release>:LINKER:--strip-all>"
)
else()
set(trt_plugins_link_options)
endif()
set(trt_plugin_dependencies
tensorrt
CUDA::cudart_static
)
if(NOT MSVC)
list(APPEND trt_plugin_dependencies
Threads::Threads
CUDA::culibos
)
endif()
set(trt_plugin_include_dirs
${TensorRT_SOURCE_DIR}/externals
${CMAKE_CURRENT_LIST_DIR}
)
set(trt_plugin_compile_options
$<$<COMPILE_LANGUAGE:CUDA>:--expt-relaxed-constexpr>
)
# Target properties for tensorrt_plugins
target_include_directories(tensorrt_plugins PRIVATE ${trt_plugin_include_dirs})
target_link_libraries(tensorrt_plugins PRIVATE ${trt_plugin_dependencies})
target_link_options(tensorrt_plugins PRIVATE ${trt_plugins_link_options})
target_compile_options(tensorrt_plugins PRIVATE ${trt_plugin_compile_options})
set_target_properties(
tensorrt_plugins
PROPERTIES CXX_VISIBILITY_PRESET hidden
VISIBILITY_INLINES_HIDDEN ON
OUTPUT_NAME nvinfer_plugin
VERSION ${TensorRT_VERSION}
SOVERSION ${TRT_MAJOR}
LINK_DEPENDS ${CMAKE_CURRENT_LIST_DIR}/exports.map)
if (NOT MSVC)
set(trt_plugins_internal_link_options
"LINKER:--version-script=${TensorRT_SOURCE_DIR}/Exports-plugin_internal.map"
"LINKER:-z,relro"
"LINKER:-Bsymbolic"
"LINKER:--no-undefined"
"LINKER:--no-as-needed"
"$<$<CONFIG:Release>:LINKER:--strip-all>"
)
else()
set(trt_plugins_internal_link_options)
endif()
# Target properties for tensorrt_plugins_internal
# This library is effectively the same as tensorrt_plugins, but without stripped symbols.
target_include_directories(tensorrt_plugins_internal PUBLIC ${trt_plugin_include_dirs})
target_link_libraries(tensorrt_plugins_internal PRIVATE ${trt_plugin_dependencies})
target_link_options(tensorrt_plugins_internal PRIVATE ${trt_plugins_internal_link_options})
target_compile_options(tensorrt_plugins_internal PRIVATE ${trt_plugin_compile_options})
set_target_properties(
tensorrt_plugins
PROPERTIES CXX_VISIBILITY_PRESET default
VISIBILITY_INLINES_HIDDEN OFF
OUTPUT_NAME nvinfer_plugin_internal
VERSION ${TensorRT_VERSION}
SOVERSION ${TRT_MAJOR}
LINK_DEPENDS ${TensorRT_SOURCE_DIR}/Exports-plugin_internal.map)
if (NOT MSVC)
set(trt_vc_plugins_link_options
"LINKER:--version-script=${CMAKE_CURRENT_LIST_DIR}/exports-vfc_plugin.map"
"LINKER:-z,relro"
"LINKER:-Bsymbolic"
"LINKER:--no-undefined"
"LINKER:--no-as-needed"
"$<$<CONFIG:Release>:LINKER:--strip-all>"
)
else()
set(trt_vc_plugins_link_options)
endif()
# Target properties for tensorrt_vc_plugins
# This library includes a minimal subset of the plugins used for version compatibility.
target_include_directories(tensorrt_vc_plugins PRIVATE ${trt_plugin_include_dirs})
target_link_libraries(tensorrt_vc_plugins PRIVATE ${trt_plugin_dependencies})
target_link_options(tensorrt_vc_plugins PRIVATE ${trt_vc_plugins_link_options})
target_compile_options(tensorrt_vc_plugins PRIVATE ${trt_plugin_compile_options})
set_target_properties(
tensorrt_vc_plugins
PROPERTIES CXX_VISIBILITY_PRESET hidden
VISIBILITY_INLINES_HIDDEN ON
OUTPUT_NAME nvinfer_vc_plugin
VERSION ${TensorRT_VERSION}
SOVERSION ${TRT_MAJOR}
LINK_DEPENDS ${CMAKE_CURRENT_LIST_DIR}/exports-vfc_plugin.map)
foreach(SM IN LISTS CMAKE_CUDA_ARCHITECTURES)
target_compile_definitions(tensorrt_plugins PRIVATE "ENABLE_SM${SM}")
target_compile_definitions(tensorrt_plugins_internal PRIVATE "ENABLE_SM${SM}")
target_compile_definitions(tensorrt_vc_plugins PRIVATE "ENABLE_SM${SM}")
endforeach()
install(
TARGETS tensorrt_plugins tensorrt_plugins_internal tensorrt_vc_plugins
OPTIONAL
LIBRARY DESTINATION "."
ARCHIVE DESTINATION "."
RUNTIME DESTINATION "."
)
else() # TRT_BUILD_ENABLE_NEW_PLUGIN_FLOW
add_custom_target(plugin)
set(TARGET_NAME ${nvinfer_plugin_lib_name})
@@ -29,14 +229,42 @@ set(VFC_PLUGIN_EXPORT_MAP ${TARGET_DIR}/exports-vfc_plugin.map)
if(${CMAKE_BUILD_TYPE} MATCHES "Debug")
set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -g")
endif()
set(ENABLED_SMS "-DENABLE_SM72 -DENABLE_SM75 -DENABLE_SM80 -DENABLE_SM86 -DENABLE_SM87 -DENABLE_SM89 -DENABLE_SM90")
if(CUDA_VERSION VERSION_GREATER_EQUAL 12.8)
set(ENABLED_SMS "${ENABLED_SMS} -DENABLE_SM100")
set(ENABLED_SMS "${ENABLED_SMS} -DENABLE_SM120")
endif()
foreach(SM IN LISTS CMAKE_CUDA_ARCHITECTURES)
set(ENABLED_SMS "${ENABLED_SMS} -DENABLE_SM${SM}")
endforeach()
set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} ${ENABLED_SMS}")
set(PLUGIN_SOURCES)
set(PLUGIN_CU_SOURCES)
set(VFC_PLUGIN_SOURCES)
set(VFC_PLUGIN_CU_SOURCES)
# Backwards compat versions of the functions used by the new flow that bind to the old flow.
# We have to add C++ sources to PLUGIN_SOURCES, and Cuda sources to PLUGIN_CU_SOURCES.
macro(add_plugin_source)
foreach(SRC_FILE IN ITEMS ${ARGN})
if(SRC_FILE MATCHES "^.*((\.h)|(\.cpp)|(\.hpp))$")
list(APPEND PLUGIN_SOURCES ${CMAKE_CURRENT_LIST_DIR}/${SRC_FILE})
else()
list(APPEND PLUGIN_CU_SOURCES ${CMAKE_CURRENT_LIST_DIR}/${SRC_FILE})
endif()
endforeach()
set(PLUGIN_SOURCES ${PLUGIN_SOURCES} PARENT_SCOPE)
set(PLUGIN_CU_SOURCES ${PLUGIN_CU_SOURCES} PARENT_SCOPE)
endmacro()
macro(add_vc_plugin_source)
foreach(SRC_FILE IN ITEMS ${ARGN})
if(SRC_FILE MATCHES "^.*((\.h)|(\.cpp)|(\.hpp))$")
list(APPEND VFC_PLUGIN_SOURCES ${CMAKE_CURRENT_LIST_DIR}/${SRC_FILE})
else()
list(APPEND VFC_PLUGIN_CU_SOURCES ${CMAKE_CURRENT_LIST_DIR}/${SRC_FILE})
endif()
endforeach()
set(VFC_PLUGIN_SOURCES ${VFC_PLUGIN_SOURCES} PARENT_SCOPE)
set(VFC_PLUGIN_CU_SOURCES ${VFC_PLUGIN_CU_SOURCES} PARENT_SCOPE)
endmacro()
set(PLUGIN_LISTS
batchTilePlugin
@@ -48,7 +276,6 @@ set(PLUGIN_LISTS
detectionLayerPlugin
disentangledAttentionPlugin
efficientNMSPlugin
efficientNMSPlugin/tftrt
flattenConcat
generateDetectionPlugin
gridAnchorPlugin
@@ -243,3 +470,5 @@ install(
RUNTIME DESTINATION bin
LIBRARY DESTINATION lib
ARCHIVE DESTINATION lib)
endif()
plugin/api/CMakeLists.txt
+18
View File
@@ -0,0 +1,18 @@
#
# SPDX-FileCopyrightText: Copyright (c) 1993-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
add_plugin_source(inferPlugin.cpp)
plugin/batchTilePlugin/CMakeLists.txt
+5 -6
View File
@@ -1,5 +1,5 @@
#
# SPDX-FileCopyrightText: Copyright (c) 1993-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-FileCopyrightText: Copyright (c) 1993-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#
# Licensed under the Apache License, Version 2.0 (the "License");
@@ -14,8 +14,7 @@
# See the License for the specific language governing permissions and
# limitations under the License.
#
file(GLOB SRCS *.cpp)
set(PLUGIN_SOURCES ${PLUGIN_SOURCES} ${SRCS})
set(PLUGIN_SOURCES
${PLUGIN_SOURCES}
PARENT_SCOPE)
add_plugin_source(
batchTilePlugin.cpp
batchTilePlugin.h
)
plugin/batchedNMSPlugin/CMakeLists.txt
+8 -11
View File
@@ -1,5 +1,5 @@
#
# SPDX-FileCopyrightText: Copyright (c) 1993-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-FileCopyrightText: Copyright (c) 1993-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#
# Licensed under the Apache License, Version 2.0 (the "License");
@@ -14,13 +14,10 @@
# See the License for the specific language governing permissions and
# limitations under the License.
#
file(GLOB SRCS *.cpp)
set(PLUGIN_SOURCES ${PLUGIN_SOURCES} ${SRCS})
set(PLUGIN_SOURCES
${PLUGIN_SOURCES}
PARENT_SCOPE)
file(GLOB CU_SRCS *.cu)
set(PLUGIN_CU_SOURCES ${PLUGIN_CU_SOURCES} ${CU_SRCS})
set(PLUGIN_CU_SOURCES
${PLUGIN_CU_SOURCES}
PARENT_SCOPE)
add_plugin_source(
batchedNMSInference.cu
batchedNMSPlugin.cpp
batchedNMSPlugin.h
gatherNMSOutputs.cu
gatherNMSOutputs.h
)
plugin/bertQKVToContextPlugin/CMakeLists.txt
+61 -11
View File
@@ -1,5 +1,5 @@
#
# SPDX-FileCopyrightText: Copyright (c) 1993-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-FileCopyrightText: Copyright (c) 1993-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#
# Licensed under the Apache License, Version 2.0 (the "License");
@@ -14,15 +14,65 @@
# See the License for the specific language governing permissions and
# limitations under the License.
#
add_plugin_source(
mhaRunner.cu
mhaRunner.h
qkvToContextInt8InterleavedPlugin.cpp
qkvToContextInt8InterleavedPlugin.h
qkvToContextInt8InterleavedPluginLegacy.cpp
qkvToContextInt8InterleavedPluginLegacy.h
qkvToContextPlugin.cpp
qkvToContextPlugin.h
qkvToContextPluginLegacy.cpp
qkvToContextPluginLegacy.h
zeroPadding2d.cu
zeroPadding2d.h
)
set(BERT_QKV_SUPPORTED_SMS
75
80
86
87
89
90
100
120
)
# Certain cubins are binary compatible between different SM versions, so they are reused.
# This function checks if a SM-named file should be compiled based on current SM enablement.
# Specifically, the SM80 files are compiled if either 80, 86, or 89 are enabled.
function(should_compile_kernel SM OUT_VAR)
# If the target SM is any of 80/86/89, we need to check if any of those are enabled in CMAKE_CUDA_ARCHITECTURES.
if((${SM} EQUAL 80) OR (${SM} EQUAL 86) OR (${SM} EQUAL 89))
list(FIND CMAKE_CUDA_ARCHITECTURES 80 SM80_INDEX)
list(FIND CMAKE_CUDA_ARCHITECTURES 86 SM86_INDEX)
list(FIND CMAKE_CUDA_ARCHITECTURES 89 SM89_INDEX)
if((NOT ${SM80_INDEX} EQUAL -1) OR
(NOT ${SM86_INDEX} EQUAL -1) OR
(NOT ${SM89_INDEX} EQUAL -1)
)
set(${OUT_VAR} TRUE PARENT_SCOPE)
else()
set(${OUT_VAR} FALSE PARENT_SCOPE)
endif()
else()
list(FIND CMAKE_CUDA_ARCHITECTURES ${SM} SM_INDEX)
if (NOT ${SM_INDEX} EQUAL -1)
set(${OUT_VAR} TRUE PARENT_SCOPE)
else()
set(${OUT_VAR} FALSE PARENT_SCOPE)
endif()
endif()
endfunction()
add_subdirectory(fused_multihead_attention)
add_subdirectory(fused_multihead_attention_v2)
file(GLOB SRCS *.cpp)
set(PLUGIN_SOURCES ${PLUGIN_SOURCES} ${SRCS})
set(PLUGIN_SOURCES
${PLUGIN_SOURCES}
PARENT_SCOPE)
file(GLOB CU_SRCS *.cu)
set(BERT_CU_SOURCES ${BERT_CU_SOURCES} ${CU_SRCS})
set(BERT_CU_SOURCES
${BERT_CU_SOURCES}
PARENT_SCOPE)
# Promote PLUGIN_SOURCES and PLUGIN_CU_SOURCES added by `add_subdirectory` to this file's parent.
if (NOT DEFINED ${TRT_BUILD_ENABLE_NEW_PLUGIN_FLOW})
set(PLUGIN_SOURCES ${PLUGIN_SOURCES} PARENT_SCOPE)
set(PLUGIN_CU_SOURCES ${PLUGIN_CU_SOURCES} PARENT_SCOPE)
endif()
plugin/bertQKVToContextPlugin/fused_multihead_attention/CMakeLists.txt
+39 -11
View File
@@ -1,5 +1,5 @@
#
# SPDX-FileCopyrightText: Copyright (c) 1993-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-FileCopyrightText: Copyright (c) 1993-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#
# Licensed under the Apache License, Version 2.0 (the "License");
@@ -14,13 +14,41 @@
# See the License for the specific language governing permissions and
# limitations under the License.
#
file(GLOB SRCS src/*.cpp)
set(PLUGIN_SOURCES ${PLUGIN_SOURCES} ${SRCS})
set(PLUGIN_SOURCES
${PLUGIN_SOURCES}
PARENT_SCOPE)
file(GLOB CU_SRCS src/*.cu)
set(BERT_CU_SOURCES ${BERT_CU_SOURCES} ${CU_SRCS})
set(BERT_CU_SOURCES
${BERT_CU_SOURCES}
PARENT_SCOPE)
# This folder contains a bunch of source files holding cubins. Since the files are huge, we only want to include and compile them if used.
# Usage is indicated by the SM being set in CMAKE_CUDA_ARCHITECTURES.
function(add_plugin_source_if_exists)
foreach(SRC_FILE IN LISTS ARGN)
if(EXISTS ${CMAKE_CURRENT_LIST_DIR}/${SRC_FILE})
add_plugin_source(${SRC_FILE})
endif()
endforeach()
endfunction()
foreach(SM IN LISTS BERT_QKV_SUPPORTED_SMS)
should_compile_kernel(${SM} SHOULD_COMPILE)
if (${SHOULD_COMPILE})
# Not every file exists for each SM, so we list all of the candidates and add them if present.
add_plugin_source_if_exists(
fused_multihead_attention_fp16_64_64_kernel.sm${SM}.cpp
fused_multihead_attention_fp16_96_64_kernel.sm${SM}.cpp
fused_multihead_attention_fp16_128_64_kernel.sm${SM}.cpp
fused_multihead_attention_fp16_256_64_kernel.sm${SM}.cpp
fused_multihead_attention_fp16_384_64_kernel.sm${SM}.cpp
fused_multihead_attention_fp16_512_64_kernel.sm${SM}.cpp
fused_multihead_attention_int8_64_64_kernel.sm${SM}.cpp
fused_multihead_attention_int8_96_64_kernel.sm${SM}.cpp
fused_multihead_attention_int8_128_64_kernel.sm${SM}.cpp
fused_multihead_attention_int8_192_64_kernel.sm${SM}.cpp
fused_multihead_attention_int8_256_64_kernel.sm${SM}.cpp
fused_multihead_attention_int8_384_64_kernel.sm${SM}.cpp
fused_multihead_attention_int8_512_64_kernel.sm${SM}.cpp
)
endif()
endforeach()
add_plugin_source(
fused_multihead_attention_common.h
)
plugin/bertQKVToContextPlugin/fused_multihead_attention/include/fused_multihead_attention.h → plugin/bertQKVToContextPlugin/fused_multihead_attention/fused_multihead_attention.h
-1
View File
@@ -478,7 +478,6 @@ extern uint32_t cubin_fmha_v1_fp16_64_64_sm120_cu_cubin_len;
|| defined(ENABLE_SM87) || defined(ENABLE_SM89) || defined(ENABLE_SM90) || defined(ENABLE_SM100) \
|| defined(ENABLE_SM120))
#error This file can only be included if one of sm 72, 75, 80, 86, 87, 89, 90, 100 or 120 is defined.
#endif
plugin/bertQKVToContextPlugin/fused_multihead_attention/include/fused_multihead_attention_common.h → plugin/bertQKVToContextPlugin/fused_multihead_attention/fused_multihead_attention_common.h
View File
plugin/bertQKVToContextPlugin/fused_multihead_attention/src/fused_multihead_attention_fp16_128_64_kernel.sm100.cpp → plugin/bertQKVToContextPlugin/fused_multihead_attention/fused_multihead_attention_fp16_128_64_kernel.sm100.cpp
View File
plugin/bertQKVToContextPlugin/fused_multihead_attention/src/fused_multihead_attention_fp16_128_64_kernel.sm120.cpp → plugin/bertQKVToContextPlugin/fused_multihead_attention/fused_multihead_attention_fp16_128_64_kernel.sm120.cpp
View File
plugin/bertQKVToContextPlugin/fused_multihead_attention/src/fused_multihead_attention_fp16_128_64_kernel.sm75.cpp → plugin/bertQKVToContextPlugin/fused_multihead_attention/fused_multihead_attention_fp16_128_64_kernel.sm75.cpp
View File
plugin/bertQKVToContextPlugin/fused_multihead_attention/src/fused_multihead_attention_fp16_128_64_kernel.sm80.cpp → plugin/bertQKVToContextPlugin/fused_multihead_attention/fused_multihead_attention_fp16_128_64_kernel.sm80.cpp
View File
plugin/bertQKVToContextPlugin/fused_multihead_attention/src/fused_multihead_attention_fp16_128_64_kernel.sm87.cpp → plugin/bertQKVToContextPlugin/fused_multihead_attention/fused_multihead_attention_fp16_128_64_kernel.sm87.cpp
View File
plugin/bertQKVToContextPlugin/fused_multihead_attention/src/fused_multihead_attention_fp16_128_64_kernel.sm90.cpp → plugin/bertQKVToContextPlugin/fused_multihead_attention/fused_multihead_attention_fp16_128_64_kernel.sm90.cpp
View File
plugin/bertQKVToContextPlugin/fused_multihead_attention/src/fused_multihead_attention_fp16_256_64_kernel.sm100.cpp → plugin/bertQKVToContextPlugin/fused_multihead_attention/fused_multihead_attention_fp16_256_64_kernel.sm100.cpp
View File
plugin/bertQKVToContextPlugin/fused_multihead_attention/src/fused_multihead_attention_fp16_256_64_kernel.sm120.cpp → plugin/bertQKVToContextPlugin/fused_multihead_attention/fused_multihead_attention_fp16_256_64_kernel.sm120.cpp
View File
plugin/bertQKVToContextPlugin/fused_multihead_attention/src/fused_multihead_attention_fp16_384_64_kernel.sm100.cpp → plugin/bertQKVToContextPlugin/fused_multihead_attention/fused_multihead_attention_fp16_384_64_kernel.sm100.cpp
View File
plugin/bertQKVToContextPlugin/fused_multihead_attention/src/fused_multihead_attention_fp16_384_64_kernel.sm120.cpp → plugin/bertQKVToContextPlugin/fused_multihead_attention/fused_multihead_attention_fp16_384_64_kernel.sm120.cpp
View File
plugin/bertQKVToContextPlugin/fused_multihead_attention/src/fused_multihead_attention_fp16_384_64_kernel.sm75.cpp → plugin/bertQKVToContextPlugin/fused_multihead_attention/fused_multihead_attention_fp16_384_64_kernel.sm75.cpp
View File
plugin/bertQKVToContextPlugin/fused_multihead_attention/src/fused_multihead_attention_fp16_384_64_kernel.sm80.cpp → plugin/bertQKVToContextPlugin/fused_multihead_attention/fused_multihead_attention_fp16_384_64_kernel.sm80.cpp
View File
plugin/bertQKVToContextPlugin/fused_multihead_attention/src/fused_multihead_attention_fp16_384_64_kernel.sm86.cpp → plugin/bertQKVToContextPlugin/fused_multihead_attention/fused_multihead_attention_fp16_384_64_kernel.sm86.cpp
View File
plugin/bertQKVToContextPlugin/fused_multihead_attention/src/fused_multihead_attention_fp16_384_64_kernel.sm87.cpp → plugin/bertQKVToContextPlugin/fused_multihead_attention/fused_multihead_attention_fp16_384_64_kernel.sm87.cpp
View File
plugin/bertQKVToContextPlugin/fused_multihead_attention/src/fused_multihead_attention_fp16_384_64_kernel.sm90.cpp → plugin/bertQKVToContextPlugin/fused_multihead_attention/fused_multihead_attention_fp16_384_64_kernel.sm90.cpp
View File

Some files were not shown because too many files have changed in this diff Show More