mirrors/nvidia--tensorrt
1
0
Fork 0
Code Issues Pull Requests Actions 11 Deployments Models Agent Notes Packages Projects Releases Wiki Activity

chore: import upstream snapshot with attribution
Docker Image CI / build-ubuntu2004 (push) Has been cancelled
Details

Browse Source
This commit is contained in:
wehub-resource-sync
2026-07-13 13:36:55 +08:00
commit c8a779b1bb
1887 changed files with 3245738 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
demo/DeBERTa/.gitignore
+5
View File
@@ -0,0 +1,5 @@
*.onnx
*.engine
*.profile
**/__pycache__/
demo/DeBERTa/README.md
+194
View File
@@ -0,0 +1,194 @@
# DeBERTa Model Inference with TensorRT Disentangled Attention Optimizations
- [DeBERTa Model Inference with TensorRT Disentangled Attention Optimizations](#deberta-model-inference-with-tensorrt-disentangled-attention-optimizations)
- [Background](#background)
- [Performance Benchmark](#performance-benchmark)
- [Environment Setup](#environment-setup)
- [Step 1: PyTorch model to ONNX model](#step-1-pytorch-model-to-onnx-model)
- [Step 2: Modify ONNX model for TensorRT engine building](#step-2-modify-onnx-model-for-tensorrt-engine-building)
- [Step 3: Model inference with TensorRT (using Python TensorRT API or `trtexec`)](#step-3-model-inference-with-tensorrt-using-python-tensorrt-api-or-trtexec)
- [Optional Step: Correctness check of model with and without plugin](#optional-step-correctness-check-of-model-with-and-without-plugin)
- [Optional Step: Model inference with ONNX Runtime and TensorRT Execution Provider (Python API)](#optional-step-model-inference-with-onnx-runtime-and-tensorrt-execution-provider-python-api)
***
## Background
A performance gap has been observed between Google's [BERT](https://arxiv.org/abs/1810.04805) design and Microsoft's [DeBERTa](https://arxiv.org/abs/2006.03654) design. The main reason of the gap is the disentangled attention design in DeBERTa triples the attention computation over BERT's regular attention. In addition to the extra matrix multiplications, the disentangled attention design also involves indirect memory accesses during the gather operations. In this regard, a [TensorRT plugin](https://github.com/NVIDIA/TensorRT/tree/main/plugin/disentangledAttentionPlugin) has been implemented to optimize DeBERTa's disentangled attention module, which is built-in since TensorRT 8.4 GA Update 2 (8.4.3) release.
This DeBERTa demo includes code and scripts for (i) exporting ONNX model from PyTorch, (ii) modifying ONNX model by inserting the plugin nodes, (iii) model inference options with TensorRT `trtexec` executable, TensorRT Python API, or ONNX Runtime with TensorRT execution provider, and (iv) measuring the correctness and performance of the optimized model.
The demo works with the [HuggingFace implementation](https://github.com/huggingface/transformers/tree/main/src/transformers/models/deberta_v2) of DeBERTa.
## Performance Benchmark
Experiments of inference performance are measured on NVIDIA A100-80GB, T4, and V100-16GB, using TensorRT 8.4 GA Update 2 (8.4.3) + CUDA 11.6 and TensorRT 8.2 GA Update 3 (8.2.4) + CUDA 11.4. The application-based model configuration is: sequence length = 512/1024/2048, hidden size = 384, intermediate size = 1536, number of layers = 12, number of attention heads = 6, maximum relative distance = 256, vocabulary size = 128K, batch size = 1, with randomly initialized weights. Numbers are average latency of 100 inference runs. Speedup numbers are fastest TensorRT number over the PyTorch baseline.
| Sequence Length | Model Latency (ms) | | TensorRT 8.4 GA Update 2 (8.4.3), CUDA 11.6 | | | TensorRT 8.2 GA Update 3 (8.2.4), CUDA 11.4 | |
| :-------------- | :----------------------------- | :-------: | :-----------------------------------------: | :-------: | :-------: | :-----------------------------------------: | :-------: |
| | | A100-80GB | T4 | V100-16GB | A100-80GB | T4 | V100-16GB |
| 512 | PyTorch (FP32/TF32) | 23.7 | 35.6 | 53.1 | 24.6 | 34.0 | 53.1 |
| | PyTorch (FP16) | 21.5 | 22.7 | 43.5 | 22.6 | 20.6 | 43.5 |
| | TensorRT (FP32/TF32) | 3.9 | 12.4 | 7.2 | 4.3 | 12.3 | 5.6 |
| | TensorRT (FP16) | **1.6** | 6.2 | **3.9** | **1.8** | 6.1 | **3.1** |
| | TensorRT w/ plugin (FP32/TF32) | 4.3 | 12.9 | 6.9 | 4.8 | 13.1 | 6.9 |
| | TensorRT w/ plugin (FP16) | 1.9 | **5.6** | 4.0 | 2.1 | **6.0** | 4.2 |
| | **Speedup** | **14.8** | **6.4** | **13.6** | **13.7** | **5.7** | **17.1** |
| | | | | | | | |
| 1024 | PyTorch (FP32/TF32) | 35.7 | 82.8 | 65.4 | 35.8 | 83.7 | 65.4 |
| | PyTorch (FP16) | 35.1 | 53.8 | 59.4 | 35.6 | 52.9 | 59.4 |
| | TensorRT (FP32/TF32) | 8.3 | 31.3 | 15.7 | 8.8 | 34.7 | 12.8 |
| | TensorRT (FP16) | 3.8 | 16.3 | 8.3 | 3.7 | 18.4 | 7.4 |
| | TensorRT w/ plugin (FP32/TF32) | 7.9 | 31.1 | 14.3 | 9.0 | 32.8 | 13.5 |
| | TensorRT w/ plugin (FP16) | **2.8** | **12.4** | **7.3** | **3.3** | **13.6** | **6.8** |
| | **Speedup** | **12.8** | **6.7** | **9.0** | **10.8** | **6.2** | **9.6** |
| | | | | | | | |
| 2048 | PyTorch (FP32/TF32) | 84.8 | 261.3 | 236.3 | 84.8 | 263.1 | 236.3 |
| | PyTorch (FP16) | 79.4 | 178.2 | 205.5 | 76.0 | 181.8 | 205.5 |
| | TensorRT (FP32/TF32) | 22.2 | 109.1 | 39.1 | 22.6 | 109.8 | 35.0 |
| | TensorRT (FP16) | 10.2 | 56.2 | 23.5 | 10.8 | 62.9 | 21.0 |
| | TensorRT w/ plugin (FP32/TF32) | 20.4 | 96.7 | 38.6 | 21.2 | 95.7 | 35.6 |
| | TensorRT w/ plugin (FP16) | **7.6** | **44.1** | **21.0** | **8.3** | **44.6** | **18.1** |
| | **Speedup** | **11.2** | **5.9** | **11.3** | **10.2** | **5.9** | **13.1** |
In addition, a pre-trained model `microsoft/deberta-v3-xsmall` was tested, which configuration is similar to sequencen length = 512 model above. And `microsoft/deberta-v3-large` (sequence length = 512) performance on TensorRT 8.4 GA Update 2 (8.4.3) is also added from recent experiments.
| Variant | Model Latency (ms) | | TensorRT 8.4 GA Update 2 (8.4.3), CUDA 11.6 | | | TensorRT 8.2 GA Update 3 (8.2.4), CUDA 11.4 | |
| :------------------ | :----------------------------- | :-------: | :-----------------------------------------: | :-------: | :-------: | :-----------------------------------------: | :-------: |
| | | A100-80GB | T4 | V100-16GB | A100-80GB | T4 | V100-16GB |
| `deberta-v3-xsmall` | PyTorch (FP32/TF32) | 30.1 | 40.6 | 57.1 | 28.9 | 39.2 | 57.1 |
| | PyTorch (FP16) | 27.6 | 26.3 | 47.9 | 26.3 | 24.8 | 47.9 |
| | TensorRT (FP32/TF32) | 4.1 | 12.4 | 7.6 | 4.4 | 12.8 | 5.8 |
| | TensorRT (FP16) | 1.8 | 6.2 | **3.7** | **1.9** | 6.4 | **3.1** |
| | TensorRT w/ plugin (FP32/TF32) | 4.3 | 12.9 | 7.7 | 4.8 | 13.6 | 6.9 |
| | TensorRT w/ plugin (FP16) | **1.8** | **5.6** | 4.7 | 2.1 | **6.0** | 4.1 |
| | **Speedup** | **16.7** | **7.3** | **15.4** | **15.2** | **6.5** | **18.4** |
| | | | | | | | |
| `deberta-v3-large` | PyTorch (FP32/TF32) | 51.6 | 40.6 | 100.0 | - | - | - |
| | PyTorch (FP16) | 52.6 | 26.3 | 96.5 | - | - | - |
| | TensorRT (FP32/TF32) | 31.1 | 32.3 | 43.2 | - | - | - |
| | TensorRT (FP16) | 7.8 | 16.0 | **13.9** | - | - | - |
| | TensorRT w/ plugin (FP32/TF32) | 30.9 | 32.8 | 44.9 | - | - | - |
| | TensorRT w/ plugin (FP16) | **7.3** | **13.5** | 14.4 | - | - | - |
| | **Speedup** | **7.1** | **3.0** | **7.2** | - | - | - |
Note that the performance gap between BERT's self-attention and DeBERTa's disentangled attention mainly comes from the additional `Gather` and `Transpose` operations in the attention design, and such gap becomes most significant when the maximum input sequence length is long (e.g., 2048). The fastest inference times are labeled as bold in the table above. In summary, for short sequence length applications, regular TensorRT inference might be sufficient, while for longer sequence length applications, the plugin optimizations are preferred and can be utilized to further improve the inference latency. Also, to get maximum speedup, using FP16 precision for inference is recommended.
## 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.8 (default)**
```bash
# Download this TensorRT OSS repo
git clone -b main https://github.com/nvidia/TensorRT TensorRT
cd TensorRT
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.8
# launch container
./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)
cd $TRT_OSSPATH
mkdir -p build && cd build
cmake .. -DTRT_LIB_DIR=$TRT_LIBPATH -DTRT_OUT_DIR=`pwd`/out
make -j$(nproc)
# polygraphy bin location & trtexec bin location
export PATH="~/.local/bin":"${TRT_OSSPATH}/build/out":$PATH
# navigate to the demo folder install additional dependencies (note PyTorch 1.11.0 is recommended for onnx to export properly)
cd $TRT_OSSPATH/demo/DeBERTa
pip install -r requirements.txt
```
> NOTE:
1. `sudo` password for Ubuntu build containers is 'nvidia'.
2. The DeBERTa plugin is only built-in after TensorRT 8.4 GA Update 2 (8.4.3) release. For pre-8.4.3 versions, you need to build TensorRT OSS from source and link the shared libraries with TensorRT build.
3. For ONNX Runtime deployment, the associated changes for the plugin are only built-in after 1.12 release. For pre-1.12 versions, you need to [build ONNX Runtime from source with TensorRT execution provider](https://onnxruntime.ai/docs/build/eps.html#tensorrt).
4. TensorRT OSS docker container is designed for use cases when you need to build OSS repository from source. After TensorRT 8.4 GA Update 2 (8.4.3) release, the most convenient way is to use the corresponding `22.08-py3` docker image at [TensorRT NGC container](https://catalog.ngc.nvidia.com/orgs/nvidia/containers/tensorrt) when it is released, without the need to build from dockerfile. But for now, please follow the instructions above to use the OSS docker image.
## Step 1: PyTorch model to ONNX model
```bash
python deberta_pytorch2onnx.py --filename ./test/deberta.onnx [--variant microsoft/deberta-v3-xsmall] [--seq-len 2048]
```
This will export the DeBERTa model from HuggingFace's DeBERTa-V2 implementation into ONNX format, with user given file name. Optionally, specify DeBERTa variant to export, such as `--variant microsoft/deberta-v3-xsmall`, or specify the maximum sequence length configuration for testing, such as `--seq-len 2048`. Models specified by `--variant` are with pre-trained weights, while models specified by `--seq-len` are with randomly initialized weights. Note that `--variant` and `--seq-len` cannot be used together because pre-trained models have pre-defined max sequence length.
## Step 2: Modify ONNX model for TensorRT engine building
```bash
python deberta_onnx_modify.py ./test/deberta.onnx # generates original TRT-compatible model, `*_original.onnx`
python deberta_onnx_modify.py ./test/deberta.onnx --plugin # generates TRT-compatible model with plugin nodes, `*_plugin.onnx`
```
The original HuggingFace implementation has uint8 Cast operations that TensorRT doesn't support, which needs to be removed from the ONNX model. After this step, the ONNX model can run in TensorRT. Without passing any flags, the script will save the model with Cast nodes removed, by default named with `_original` suffix.
Further, to use the DeBERTa plugin optimizations, the disentangled attention subgraph needs to be replaced by node named `DisentangledAttention_TRT`. By passing `--plugin` flag, the script will save the model with Cast nodes removed and plugin nodes replaced, by default named with `_plugin` suffix.
The benefits of the DeBERTa plugin optimizations can be demonstrated by comparing the latency of original model and plugin model.
## Step 3: Model inference with TensorRT (using Python TensorRT API or `trtexec`)
```bash
# build and test the original DeBERTa model (baseline)
python deberta_tensorrt_inference.py --onnx=./test/deberta_original.onnx --build fp16 --test fp16
# build and test the optimized DeBERTa model with plugin
python deberta_tensorrt_inference.py --onnx=./test/deberta_plugin.onnx --build fp16 --test fp16
```
This will build and test the original and optimized DeBERTa models. `--build` to build the engine from ONNX model, and `--test` to measure the optimized latency. TensorRT engine of different precisions (`--fp32/--tf32/--fp16`) can be built. Engine files are saved as `**/[Model name]_[GPU name]_[Precision].engine`.
> NOTE:
1. To get maximum speedup, using `--fp16` is recommended. Also, it was observed the plugin optimizations demonstrate more speedup when the input sequence length is long, such as 2048.
2. TF32 is only effective on Ampere and later devices.
3. TensorRT optimization profile in `deberta_tensorrt_inference.py` is set by default for batch size of 1 usage. For batch usage, the optional and maximum profile should be changed.
4. TensorRT engines are specific to the exact GPU device they were built on, as well as the platform and the TensorRT version. On the same machine, building is needed only once and the engine can be used for repeated testing.
For `trtexec` test, it is recommended to used the [TensorRT NGC container](https://catalog.ngc.nvidia.com/orgs/nvidia/containers/tensorrt) where the executable is built-in. The DeBERTa plugin support is available since docker image `22.08-py3`, which contains TensorRT 8.4 GA Update 2 (8.4.3) pre-installed. For earlier images before `22.08`, skip this `trtexec` test.
`trtexec` command for sequence length = 2048 model is given as an example:
```bash
trtexec --onnx=./test/deberta_plugin.onnx --workspace=4096 --explicitBatch --optShapes=input_ids:1x2048,attention_mask:1x2048 --iterations=10 --warmUp=10 --noDataTransfers --fp16
```
## Optional Step: Correctness check of model with and without plugin
```bash
# prepare the ONNX models with intermediate output nodes (this will save two new onnx models with suffix `*_correctness_check_original.onnx` and `*_correctness_check_plugin.onnx`)
python deberta_onnx_modify.py ./test/deberta.onnx --correctness-check
# build the ONNX models with intermediate outputs for comparison
python deberta_tensorrt_inference.py --onnx=./test/deberta_correctness_check_original.onnx --build fp16
python deberta_tensorrt_inference.py --onnx=./test/deberta_correctness_check_plugin.onnx --build fp16
# run correctness check (specify the root model name with --onnx)
python deberta_tensorrt_inference.py --onnx=./test/deberta --correctness-check fp16
```
Correctness check requires intermediate outputs from the model, thus it is necessary to modify the ONNX graph and add intermediate output nodes. The correctness check was added at the location of plugin outputs in each layer. The metrics are average and maximum of the element-wise absolute error. Note that for FP16 precision with 10 significance bits, absolute error in the order of 1e-2 and 1e-3 is expected, and for FP32 precision with 23 significance bits, 1e-6 to 1e-7 is expected. Refer to [Machine Epsilon](https://en.wikipedia.org/wiki/Machine_epsilon) for details.
Example output for FP16 correctness check between original DeBERTa model and plugin optimized DeBERTa model:
```bash
[Layer 0 Element-wise Check] Avgerage absolute error: 2.152580e-03, Maximum absolute error: 3.010559e-02. 1e-2~1e-3 expected for FP16 (10 significance bits) and 1e-6~1e-7 expected for FP32 (23 significance bits)
...
...
[Layer 12 Element-wise Check] Avgerage absolute error: 6.198883e-05, Maximum absolute error: 1.220703e-04. 1e-2~1e-3 expected for FP16 (10 significance bits) and 1e-6~1e-7 expected for FP32 (23 significance bits)
```
## Optional Step: Model inference with ONNX Runtime and TensorRT Execution Provider (Python API)
```bash
python deberta_ort_inference.py --onnx=./test/deberta_original.onnx --test fp16
python deberta_ort_inference.py --onnx=./test/deberta_plugin.onnx --test fp16
python deberta_ort_inference.py --onnx=./test/deberta --correctness-check fp16 # for correctness check
```
In addition to TensorRT inference, [ONNX Runtime](https://github.com/microsoft/onnxruntime) with TensorRT Execution Provider can also be used as the inference framework. The DeBERTa TensorRT plugin is officially supported in onnxruntime since version 1.12. For earlier releases of onnxruntime, skip this step.
The results can be cross-validated between TensorRT and onnxruntime. For example, the [Polygraphy](https://github.com/NVIDIA/TensorRT/tree/master/tools/Polygraphy) tool can be used for easy comparison:
```bash
polygraphy run ./test/deberta_original.onnx --trt --onnxrt --workspace=4000000000
```
demo/DeBERTa/cuda_utils.py
+128
View File
@@ -0,0 +1,128 @@
#!/usr/bin/env python3
#
# 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 numpy as np
import logging
from cuda.bindings import driver as cuda, runtime as cudart
class CudaStreamContext:
"""CUDA stream lifecycle management with context manager support"""
def __init__(self):
"""Initialize CUDA stream"""
self._stream = cuda_call(cudart.cudaStreamCreate())
def __enter__(self):
"""Create CUDA stream when entering context (if not already created)"""
if self._stream is None:
self._stream = cuda_call(cudart.cudaStreamCreate())
return self
def __exit__(self, exc_type, exc_val, exc_tb):
"""Destroy CUDA stream when exiting context"""
self.free()
@property
def stream(self) -> cudart.cudaStream_t:
if self._stream is None:
raise RuntimeError("Stream not created. Use 'with' statement.")
return self._stream
def synchronize(self):
"""Synchronize the stream"""
if self._stream is None:
raise RuntimeError("Stream not created. Use 'with' statement.")
cuda_call(cudart.cudaStreamSynchronize(self._stream))
def free(self):
"""Explicitly free the CUDA stream"""
if self._stream is not None:
try:
cuda_call(cudart.cudaStreamDestroy(self._stream))
self._stream = None
except Exception as e:
logging.warning(f"Failed to destroy CUDA stream: {e}")
def __del__(self):
"""Cleanup stream on destruction"""
if hasattr(self, '_stream') and self._stream is not None:
self.free()
def __str__(self):
return f"CudaStreamContext: {self._stream}"
def __repr__(self):
return self.__str__()
def cuda_call(call):
"""Helper function to make CUDA calls and check for errors"""
def _cudaGetErrorEnum(error):
if isinstance(error, cuda.CUresult):
err, name = cuda.cuGetErrorName(error)
return name if err == cuda.CUresult.CUDA_SUCCESS else "<unknown>"
elif isinstance(error, cudart.cudaError_t):
return cudart.cudaGetErrorName(error)[1]
else:
raise RuntimeError("Unknown error type: {}".format(error))
err, res = call[0], call[1:]
if err.value:
raise RuntimeError(
"CUDA error code={}({})".format(
err.value, _cudaGetErrorEnum(err)
)
)
if len(res) == 1:
return res[0]
elif len(res) == 0:
return None
else:
return res
def getComputeCapacity(devID=0):
major = cuda_call(cudart.cudaDeviceGetAttribute(cudart.cudaDeviceAttr.cudaDevAttrComputeCapabilityMajor, devID))
minor = cuda_call(cudart.cudaDeviceGetAttribute(cudart.cudaDeviceAttr.cudaDevAttrComputeCapabilityMinor, devID))
return (major, minor)
def memcpy_host_to_device_async(device_ptr: int, host_arr: np.ndarray, stream):
"""Wrapper for async host-to-device memory copy"""
cuda_call(cudart.cudaMemcpyAsync(device_ptr, host_arr.ctypes.data, host_arr.nbytes, cudart.cudaMemcpyKind.cudaMemcpyHostToDevice, stream))
def memcpy_device_to_host_async(host_arr: np.ndarray, device_ptr: int, stream):
"""Wrapper for async device-to-host memory copy"""
cuda_call(cudart.cudaMemcpyAsync(host_arr.ctypes.data, device_ptr, host_arr.nbytes, cudart.cudaMemcpyKind.cudaMemcpyDeviceToHost, stream))
def memcpy_device_to_device_async(dst_device_ptr: int, src_device_ptr: int, nbytes: int, stream):
"""Wrapper for async device-to-device memory copy"""
cuda_call(cudart.cudaMemcpyAsync(dst_device_ptr, src_device_ptr, nbytes, cudart.cudaMemcpyKind.cudaMemcpyDeviceToDevice, stream))
def memcpy_host_to_device(device_ptr: int, host_arr: np.ndarray):
"""Wrapper for synchronous host-to-device memory copy"""
cuda_call(cudart.cudaMemcpy(device_ptr, host_arr.ctypes.data, host_arr.nbytes, cudart.cudaMemcpyKind.cudaMemcpyHostToDevice))
def memcpy_device_to_host(host_arr: np.ndarray, device_ptr: int):
"""Wrapper for synchronous device-to-host memory copy"""
cuda_call(cudart.cudaMemcpy(host_arr.ctypes.data, device_ptr, host_arr.nbytes, cudart.cudaMemcpyKind.cudaMemcpyDeviceToHost))
def memcpy_device_to_device(dst_device_ptr: int, src_device_ptr: int, nbytes: int):
"""Wrapper for synchronous device-to-device memory copy"""
cuda_call(cudart.cudaMemcpy(dst_device_ptr, src_device_ptr, nbytes, cudart.cudaMemcpyKind.cudaMemcpyDeviceToDevice))
# Initialize CUDA
cuda_call(cudart.cudaFree(0))
demo/DeBERTa/deberta_onnx_modify.py
+227
View File
@@ -0,0 +1,227 @@
#
# 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.
#
'''
Modify original ONNX exported from HuggingFace to for TensorRT engine building.
The original HuggingFace implementation has uint8 Cast operations that TensorRT doesn't support, which needs to be removed from the ONNX model. After this step, the ONNX model can run in TensorRT.
Further, to use the DeBERTa plugin optimizations, the disentangled attention module needs to be replaced by node named `DisentangledAttention_TRT`.
Optional: generate model that has per-layer intermediate outputs for correctness check purpose.
These modifications are automated in this script.
Usage:
python deberta_onnx_modify.py xx.onnx # for original TRT-compatible model, `xx_original.onnx`
python deberta_onnx_modify.py xx.onnx --plugin # for TRT-compatible model with plugin nodes, `xx_plugin.onnx`
python deberta_onnx_modify.py xx.onnx --correctness-check # for correctness check
'''
import onnx
from onnx import TensorProto
import onnx_graphsurgeon as gs
import argparse, os
import numpy as np
parser = argparse.ArgumentParser(description="Modify DeBERTa ONNX model to prepare for TensorRT engine building. If none of --plugin or --correctness-check flag is passed, it will just save the uint8 cast removed model.")
parser.add_argument('input', type=str, help='Path to the input ONNX model')
parser.add_argument('--plugin', action='store_true', help="Generate model with plugin")
parser.add_argument('--correctness-check', action='store_true', help="Generate model that has per-layer intermediate outputs for correctness check purpose")
parser.add_argument('--output', type=str, help="Path to the output ONNX model. If not set, default to the input file name with a suffix of '_original' or '_plugin' ")
args = parser.parse_args()
model_input = args.input
use_plugin = args.plugin
correctness_check = args.correctness_check
if args.output is None:
model_output = os.path.splitext(model_input)[0] + ("_plugin" if use_plugin else "_original") + os.path.splitext(model_input)[-1]
else:
model_output = args.output
def remove_uint8_cast(graph):
'''
Remove all uint8 Cast nodes since TRT doesn't support UINT8 cast op.
Ref: https://github.com/NVIDIA/TensorRT/tree/main/tools/onnx-graphsurgeon/examples/06_removing_nodes
'''
nodes = [node for node in graph.nodes if node.op == 'Cast' and node.attrs["to"] == TensorProto.UINT8] # find by op name and attribute
for node in nodes:
# [ONNX's Cast operator](https://github.com/onnx/onnx/blob/main/docs/Operators.md#Cast) will exactly have 1 input and 1 output
# reconnect tensors
input_node = node.i()
input_node.outputs = node.outputs
node.outputs.clear()
# an alternative way is to just not cast to uint8
# node.attrs["to"] = TensorProto.INT64
return graph
@gs.Graph.register()
def insert_disentangled_attention(self, inputs, outputs, factor, span):
'''
Fuse disentangled attention module (Add + Gather + Gather + Transpose + Add + Div)
inputs: list of plugin inputs
outputs: list of plugin outputs
factor: scaling factor of disentangled attention, sqrt(3d), converted from a division factor to a multiplying factor
span: relative distance span, k
'''
# disconnect previous output from flow (the previous subgraph still exists but is effectively dead since it has no link to an output tensor, and thus will be cleaned up)
[out.inputs.clear() for out in outputs]
# add plugin layer
attrs = {
"factor": 1/factor,
"span": span
}
self.layer(op='DisentangledAttention_TRT', inputs=inputs, outputs=outputs, attrs=attrs)
def insert_disentangled_attention_all(graph):
'''
Insert disentangled attention plugin nodes for all layers
'''
nodes = [node for node in graph.nodes if node.op == 'GatherElements'] # find entry points by gatherelements op
assert len(nodes) % 2 == 0, "No. of GatherElements nodes is not an even number!"
layers = [(nodes[2*i+0], nodes[2*i+1]) for i in range(len(nodes)//2)] # 2 gatherelements in 1 layer
for l, (left,right) in enumerate(layers):
print(f"Fusing layer {l}")
# CAVEAT! MUST cast to list() when setting the inputs & outputs. graphsurgeon's default for X.inputs and X.outputs is `onnx_graphsurgeon.util.misc.SynchronizedList`, i.e. 2-way node-tensor updating mechanism. If not cast, when we remove the input nodes of a tensor, the tensor itself will be removed as well...
# inputs: (data0, data1, data2), input tensors for c2c add and 2 gathers
inputs = list(left.o().o().o().o().i().inputs)[0:1] + list(left.inputs)[0:1] + list(right.inputs)[0:1]
# outputs: (result), output tensors after adding 3 gather results
outputs = list(left.o().o().o().o().outputs)
# constants: scaling factor, relative distance span
factor = left.o().inputs[1].inputs[0].attrs["value"].values.item()
span = right.i(1,0).i().i().i().inputs[1].inputs[0].attrs["value"].values.item()
# insert plugin layer
graph.insert_disentangled_attention(inputs, outputs, factor, span)
return graph
def correctness_check_models(graph):
'''
Add output nodes at the plugin exit point for both the original model and the model with plugin
'''
seq_len = graph.inputs[0].shape[1]
## for original graph
# make a copy of the graph first
graph_raw = graph.copy()
nodes = [node for node in graph_raw.nodes if node.op == 'GatherElements'] # find by gatherelements op
assert len(nodes) % 2 == 0, "No. of GatherElements nodes is not an even number!"
layers = [(nodes[2*i+0], nodes[2*i+1]) for i in range(len(nodes)//2)] # 2 gatherelements in 1 layer
original_output_all = []
for l, (left,right) in enumerate(layers):
# outputs: (result), output tensors after adding 3 gather results
# add the output tensor to the graph outputs list. Don't create any new tensor!
end_node = left.o().o().o().o()
end_node.outputs[0].dtype = graph_raw.outputs[0].dtype # need to explicitly specify dtype and shape of graph output tensor
end_node.outputs[0].shape = ['batch_size*num_heads', seq_len, seq_len]
original_output_all.append(end_node.outputs[0])
graph_raw.outputs = graph_raw.outputs + original_output_all # add plugin outputs to graph output
## for modified graph with plugin
nodes = [node for node in graph.nodes if node.op == 'GatherElements'] # find by gatherelements op
assert len(nodes) % 2 == 0, "No. of GatherElements nodes is not an even number!"
layers = [(nodes[2*i+0], nodes[2*i+1]) for i in range(len(nodes)//2)] # 2 gatherelements in 1 layer
plugin_output_all = []
for l, (left,right) in enumerate(layers):
# inputs: (data0, data1, data2), input tensors for c2c add and 2 gathers
inputs = list(left.o().o().o().o().i().inputs)[0:1] + list(left.inputs)[0:1] + list(right.inputs)[0:1]
# outputs: (result), output tensors after adding 3 gather results
outputs = list(left.o().o().o().o().outputs)
end_node = left.o().o().o().o()
end_node.outputs[0].dtype = graph.outputs[0].dtype # need to explicitly specify dtype and shape of graph output tensor
end_node.outputs[0].shape = ['batch_size*num_heads', seq_len, seq_len]
plugin_output_all.append(end_node.outputs[0]) # add to graph output (outside this loop)
# constants: scaling factor, relative distance span
factor = left.o().inputs[1].inputs[0].attrs["value"].values.item()
span = right.i(1,0).i().i().i().inputs[1].inputs[0].attrs["value"].values.item()
# insert plugin layer
graph.insert_disentangled_attention(inputs, outputs, factor, span)
graph.outputs = graph.outputs + plugin_output_all # add plugin outputs to graph output
return graph_raw, graph
def check_model(model_name):
# Load the ONNX model
model = onnx.load(model_name)
# Check that the model is well formed
onnx.checker.check_model(model)
# load onnx
graph = gs.import_onnx(onnx.load(model_input))
# first, remove uint8 cast nodes
graph = remove_uint8_cast(graph)
if use_plugin:
# save the modified model with plugin nodes
# replace Add + Gather + Gather + Transpose + Add + Div (c2c and c2p and p2c) with DisentangledAttention_TRT node
graph = insert_disentangled_attention_all(graph)
# remove unused nodes, and topologically sort the graph.
graph.cleanup().toposort()
# export the onnx graph from graphsurgeon
onnx.save_model(gs.export_onnx(graph), model_output)
print(f"Saving modified model to {model_output}")
# don't check model because 'DisentangledAttention_TRT' is not a registered op
elif correctness_check:
# correctness check, save two models (original and w/ plugin) with intermediate output nodes inserted
graph_raw, graph = correctness_check_models(graph)
# remove unused nodes, and topologically sort the graph.
graph_raw.cleanup().toposort()
graph.cleanup().toposort()
# export the onnx graph from graphsurgeon
model_output1 = os.path.splitext(model_input)[0] + "_correctness_check_original" + os.path.splitext(model_input)[-1]
model_output2 = os.path.splitext(model_input)[0] + "_correctness_check_plugin" + os.path.splitext(model_input)[-1]
onnx.save_model(gs.export_onnx(graph_raw), model_output1)
onnx.save_model(gs.export_onnx(graph), model_output2)
print(f"Saving models for correctness check to {model_output1} (original) and {model_output2} (with plugin)")
check_model(model_output1)
# don't check model_output2 because 'DisentangledAttention_TRT' is not a registered op
else:
# no flag passed, save model with just uint8 cast removed
graph.cleanup().toposort()
onnx.save_model(gs.export_onnx(graph), model_output)
print(f"Saving modified model to {model_output}")
check_model(model_output)
demo/DeBERTa/deberta_ort_inference.py
+201
View File
@@ -0,0 +1,201 @@
#
# 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.
#
"""
Test ORT-TRT engine of DeBERTa model. Different precisions are supported.
Usage:
Test model inference time:
- python deberta_ort_inference.py --onnx=./test/deberta.onnx --test fp16
Correctness check by comparing original model and model with plugin:
- python deberta_ort_inference.py --onnx=./test/deberta --correctness-check fp16
Notes:
- supported precisions are fp32/fp16. For test, you can specify more than one precisions, and TensorRT engine of each precision will be built sequentially.
- engine files are saved at `./engine_cache/[Model name]_[GPU name]_[Precision]/`. Note that TensorRT engine is specific to both GPU architecture and TensorRT version.
- if in --correctness-check mode, the argument for --onnx is the stem name for the model without .onnx extension.
"""
import os, argparse
import onnxruntime as ort
import numpy as np
import torch
from time import time
ENGINE_PATH = './test'
if not os.path.exists(ENGINE_PATH):
os.makedirs(ENGINE_PATH)
def GPU_ABBREV(name):
'''
Map GPU device query name to abbreviation.
::param str name Device name from torch.cuda.get_device_name().
::return str GPU abbreviation.
'''
GPU_LIST = [
'V100',
'TITAN',
'T4',
'A100',
'A10G',
'A10'
]
# Partial list, can be extended. The order of A100, A10G, A10 matters. They're put in a way to not detect substring A10 as A100
for i in GPU_LIST:
if i in name:
return i
return 'GPU' # for names not in the partial list, use 'GPU' as default
gpu_name = GPU_ABBREV(torch.cuda.get_device_name())
VALID_PRECISION = [
'fp32',
'fp16',
]
parser = argparse.ArgumentParser(description="Build and test TensorRT engine.")
parser.add_argument('--onnx', required=True, help='ONNX model path (or filename stem if in correctness check mode).')
parser.add_argument('--test', nargs='+', help='Test ORT-TRT engine in precision fp32/fp16. You can list multiple precisions to test all of them.')
parser.add_argument('--correctness-check', nargs='+', help='Correctness check for original & plugin TRT engines in precision fp32/fp16. You can list multiple precisions to check all of them.')
args = parser.parse_args()
ONNX_MODEL = args.onnx
MODEL_STEM = os.path.splitext(args.onnx)[0].split('/')[-1]
TEST = args.test
CORRECTNESS = args.correctness_check
if TEST:
for i in TEST:
if i not in VALID_PRECISION:
parser.error(f'Unsupported precision {i}')
if CORRECTNESS:
for i in CORRECTNESS:
if i not in VALID_PRECISION:
parser.error(f'Unsupported precision {i}')
def test_engine():
for precision in TEST:
engine_cachepath = '/'.join([ENGINE_PATH, '_'.join([MODEL_STEM, gpu_name, precision, 'ort'])])
providers = [
('TensorrtExecutionProvider', {
'trt_max_workspace_size': 2147483648,
'trt_fp16_enable': precision == 'fp16',
'trt_engine_cache_enable': True,
'trt_engine_cache_path': engine_cachepath
}),
'CUDAExecutionProvider'] # EP order indicates priority
so = ort.SessionOptions()
sess = ort.InferenceSession(ONNX_MODEL, sess_options=so, providers=providers)
print(f'Running inference on engine {engine_cachepath}')
## psuedo-random input test
batch_size = 1
seq_len = sess.get_inputs()[0].shape[1]
vocab = 128203
input_ids = torch.randint(0, vocab, (batch_size, seq_len), dtype=torch.long)
attention_mask = torch.randint(0, 2, (batch_size, seq_len), dtype=torch.long)
inputs = {
'input_ids': input_ids.numpy(),
'attention_mask': attention_mask.numpy()
}
outputs = sess.run(None, inputs)
nreps = 100
start_time = time()
for _ in range(nreps):
sess.run(None, inputs)
end_time = time()
duration = end_time - start_time
print(f'Average Inference time (ms) of {nreps} runs: {duration/nreps*1000:.3f}. For more accurate test, please use the onnxruntime_perf_test commands.')
def correctness_check_engines():
for precision in CORRECTNESS:
engine_cachepath1 = '/'.join([ENGINE_PATH, '_'.join([MODEL_STEM, 'original', gpu_name, precision, 'ort'])])
engine_cachepath2 = '/'.join([ENGINE_PATH, '_'.join([MODEL_STEM, 'plugin', gpu_name, precision, 'ort'])])
if not os.path.exists(engine_cachepath1) or not os.path.exists(engine_cachepath2):
print('At least one of the original and/or plugin engines do not exist. Please build them first by --test')
return
print(f'Running inference on original engine {engine_cachepath1} and plugin engine {engine_cachepath2}')
so = ort.SessionOptions()
providers1 = [
('TensorrtExecutionProvider', {
'trt_max_workspace_size': 2147483648,
'trt_fp16_enable': precision == 'fp16',
'trt_engine_cache_enable': True,
'trt_engine_cache_path': engine_cachepath1
}),
'CUDAExecutionProvider']
providers2 = [
('TensorrtExecutionProvider', {
'trt_max_workspace_size': 2147483648,
'trt_fp16_enable': precision == 'fp16',
'trt_engine_cache_enable': True,
'trt_engine_cache_path': engine_cachepath2
}),
'CUDAExecutionProvider']
sess1 = ort.InferenceSession(ONNX_MODEL+'_original.onnx', sess_options=so, providers=providers1)
sess2 = ort.InferenceSession(ONNX_MODEL+'_plugin.onnx', sess_options=so, providers=providers2)
## psuedo-random input test
batch_size = 1
seq_len = sess1.get_inputs()[0].shape[1]
vocab = 128203
input_ids = torch.randint(0, vocab, (batch_size, seq_len), dtype=torch.long)
attention_mask = torch.randint(0, 2, (batch_size, seq_len), dtype=torch.long)
inputs = {
'input_ids': input_ids.numpy(),
'attention_mask': attention_mask.numpy()
}
outputs1 = sess1.run(None, inputs)
outputs2 = sess2.run(None, inputs)
for i in range(len(outputs1)):
avg_abs_error = np.sum(np.abs(outputs1[i] - outputs2[i])) / outputs1[i].size
max_abs_error = np.max(np.abs(outputs1[i] - outputs2[i]))
print(f"Output {i}:")
print("onnx model (original): ", outputs1[i])
print("onnx model (plugin): ", outputs2[i])
print(f"[Output {i} Element-wise Check] Avgerage absolute error: {avg_abs_error:e}, Maximum absolute error: {max_abs_error:e}. 1e-2~1e-3 expected for FP16 (10 significance bits) and 1e-6~1e-7 expected for FP32 (23 significance bits) " ) # machine epsilon for different precisions: https://en.wikipedia.org/wiki/Machine_epsilon
if TEST:
test_engine()
if CORRECTNESS:
correctness_check_engines()
demo/DeBERTa/deberta_pytorch2onnx.py
+134
View File
@@ -0,0 +1,134 @@
#
# 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.
#
'''
Generate HuggingFace DeBERTa (V2) model with different configurations (e.g., sequence length, hidden size, No. of layers, No. of heads, etc.) and export in ONNX format
Usage:
python deberta_pytorch2onnx.py [--filename xx.onnx] [--variant microsoft/deberta-xx] [--seq-len xx]
'''
import os, time, argparse
from transformers import DebertaV2Tokenizer, DebertaV2Config, DebertaV2ForSequenceClassification
# DEBERTA V2 implementation, https://github.com/huggingface/transformers/blob/master/src/transformers/models/deberta_v2/modeling_deberta_v2.py
import torch, onnxruntime as ort, numpy as np
parser = argparse.ArgumentParser(description="Generate HuggingFace DeBERTa (V2) model with different configurations and export in ONNX format. This will save the model under the same directory as 'deberta_seqxxx_hf.onnx'.")
parser.add_argument('--filename', type=str, help='Path to the save the ONNX model')
parser.add_argument('--variant', type=str, default=None, help='DeBERTa variant name. Such as microsoft/deberta-v3-xsmall')
parser.add_argument('--seq-len', type=int, default=None, help='Specify maximum sequence length. Note: --variant and --seq-len cannot be used together. Pre-trained models have pre-defined sequence length')
args = parser.parse_args()
onnx_filename = args.filename
model_variant = args.variant
sequence_length = args.seq_len
assert not args.variant or (args.variant and not args.seq_len), "--variant and --seq-len cannot be used together!"
assert torch.cuda.is_available(), "CUDA not available!"
def randomize_model(model):
for module_ in model.named_modules():
if isinstance(module_[1],(torch.nn.Linear, torch.nn.Embedding)):
module_[1].weight.data.normal_(mean=0.0, std=model.config.initializer_range)
elif isinstance(module_[1], torch.nn.LayerNorm):
module_[1].bias.data.zero_()
module_[1].weight.data.fill_(1.0)
if isinstance(module_[1], torch.nn.Linear) and module_[1].bias is not None:
module_[1].bias.data.zero_()
return model
def export():
parent_dir = os.path.dirname(onnx_filename)
if not os.path.exists(parent_dir):
os.makedirs(parent_dir)
if model_variant is None:
# default model hyper-params
batch_size = 1
seq_len = 2048 if sequence_length is None else sequence_length
max_position_embeddings = 512 if seq_len <= 512 else seq_len # maximum sequence length that this model might ever be used with. By default 512. otherwise error https://github.com/huggingface/transformers/issues/4542
vocab_size = 128203
hidden_size = 384
layers = 12
heads = 6
intermediate_size = hidden_size*4 # feed forward layer dimension
type_vocab_size = 0
# relative attention
relative_attention=True
max_relative_positions = 256 # k
pos_att_type = ["p2c", "c2p"]
deberta_config = DebertaV2Config(vocab_size=vocab_size, hidden_size=hidden_size, num_hidden_layers=layers, num_attention_heads=heads, intermediate_size=intermediate_size, type_vocab_size=type_vocab_size, max_position_embeddings=max_position_embeddings, relative_attention=relative_attention, max_relative_positions=max_relative_positions, pos_att_type=pos_att_type)
deberta_model = DebertaV2ForSequenceClassification(deberta_config)
deberta_model = randomize_model(deberta_model)
else:
deberta_model = DebertaV2ForSequenceClassification.from_pretrained(model_variant)
deberta_config = DebertaV2Config.from_pretrained(model_variant)
batch_size = 1
seq_len = deberta_config.max_position_embeddings
vocab_size = deberta_config.vocab_size
deberta_model.cuda().eval()
# input/output
gpu = torch.device('cuda')
input_ids = torch.randint(0, vocab_size, (batch_size, seq_len), dtype=torch.long, device=gpu)
attention_mask = torch.randint(0, 2, (batch_size, seq_len), dtype=torch.long, device=gpu)
input_names = ['input_ids', 'attention_mask']
output_names = ['output']
dynamic_axes={'input_ids' : {0 : 'batch_size'},
'attention_mask' : {0 : 'batch_size'},
'output' : {0 : 'batch_size'}}
# ONNX export
torch.onnx.export(deberta_model, # model
(input_ids, attention_mask), # model inputs
onnx_filename,
export_params=True,
opset_version=13,
do_constant_folding=True,
input_names = input_names,
output_names = output_names,
dynamic_axes = dynamic_axes)
# full precision inference
num_trials = 10
start = time.time()
for i in range(num_trials):
results = deberta_model(input_ids, attention_mask)
end = time.time()
print("Average PyTorch FP32(TF32) time: {:.2f} ms".format((end - start)/num_trials*1000))
# half precision inference (do this after onnx export, otherwise the export ONNX model is with FP16 weights...)
deberta_model_fp16 = deberta_model.half()
start = time.time()
for i in range(num_trials):
results = deberta_model_fp16(input_ids, attention_mask)
end = time.time()
print("Average PyTorch FP16 time: {:.2f} ms".format((end - start)/num_trials*1000))
# model size
total_params = sum(param.numel() for param in deberta_model.parameters())
print("Total # of params: ", total_params)
print("Maximum sequence length: ", seq_len)
if __name__ == "__main__":
export()
demo/DeBERTa/deberta_tensorrt_inference.py
+401
View File
@@ -0,0 +1,401 @@
#
# 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.
#
"""
Build and test TensorRT engines generated from the DeBERTa model. Different precisions are supported.
Usage:
Build and test a model:
- build: python deberta_tensorrt_inference.py --onnx=xx.onnx --build fp16 # build TRT engines
- test: python deberta_tensorrt_inference.py --onnx=xx.onnx --test fp16 # test will measure the inference time
- build and test: python deberta_tensorrt_inference.py --onnx=xx.onnx --build fp16 --test fp16
Correctness check is done by comparing engines generated from the original model and the plugin model:
- [1] export ONNX model with extra output nodes: python deberta_onnx_modify.py xx.onnx --correctness-check
- [2] build original model: python deberta_tensorrt_inference.py --onnx=xx_correctness_check_original.onnx --build fp16
- [3] build plugin model: python deberta_tensorrt_inference.py --onnx=xx_correctness_check_plugin.onnx --build fp16
- [4] correctness check: python deberta_tensorrt_inference.py --onnx=deberta --correctness_check fp16
Notes:
- supported precisions are fp32/tf32/fp16. For both --build and --test, you can specify more than one precisions, and TensorRT engines of each precision will be built sequentially.
- engine files are saved as `**/[Model name]_[GPU name]_[Precision].engine`. Note that TensorRT engines are specific to both GPU architecture and TensorRT version, and therefore are not compatible cross-version nor cross-device.
- in --correctness-check mode, the argument for --onnx is the `root` name for the models [root]_correctness_check_original/plugin.onnx
"""
import torch
import tensorrt as trt
import os, sys, argparse
import numpy as np
from time import time
from cuda.bindings import driver as cuda, runtime as cudart
from cuda_utils import (
cuda_call,
CudaStreamContext,
memcpy_host_to_device_async,
memcpy_device_to_host_async,
memcpy_host_to_device,
memcpy_device_to_device_async,
memcpy_device_to_device,
)
TRT_VERSION = int(trt.__version__[:3].replace('.','')) # e.g., version 8.4.1.5 becomes 84
def GPU_ABBREV(name):
'''
Map GPU device query name to abbreviation.
::param str name Device name from torch.cuda.get_device_name().
::return str GPU abbreviation.
'''
GPU_LIST = [
'V100',
'TITAN',
'T4',
'A100',
'A10G',
'A10'
]
# Partial list, can be extended. The order of A100, A10G, A10 matters. They're put in a way to not detect substring A10 as A100
for i in GPU_LIST:
if i in name:
return i
return 'GPU' # for names not in the partial list, use 'GPU' as default
gpu_name = GPU_ABBREV(torch.cuda.get_device_name())
VALID_PRECISION = [
'fp32',
'tf32',
'fp16'
]
parser = argparse.ArgumentParser(description="Build and test TensorRT engine.")
parser.add_argument('--onnx', required=True, help='ONNX model path (or filename stem if in correctness check mode).')
parser.add_argument('--build', nargs='+', help='Build TRT engine in precision fp32/tf32/fp16. You can list multiple precisions to build all of them.')
parser.add_argument('--test', nargs='+', help='Test TRT engine in precision fp32/tf32/fp16. You can list multiple precisions to test all of them.')
parser.add_argument('--correctness-check', nargs='+', help='Correctness check for original & plugin TRT engines in precision fp32/tf32/fp16. You can list multiple precisions to check all of them.')
args = parser.parse_args()
ONNX_MODEL = args.onnx
MODEL_NAME = os.path.splitext(args.onnx)[0]
BUILD = args.build
TEST = args.test
CORRECTNESS = args.correctness_check
if not (args.build or args.test or args.correctness_check):
parser.error('Please specify --build and/or --test and/or --correctness-check' )
if BUILD:
for i in BUILD:
if i not in VALID_PRECISION:
parser.error(f'Unsupported precision {i}')
if TEST:
for i in TEST:
if i not in VALID_PRECISION:
parser.error(f'Unsupported precision {i}')
if CORRECTNESS:
for i in CORRECTNESS:
if i not in VALID_PRECISION:
parser.error(f'Unsupported precision {i}')
class TRTModel:
'''
Generic class to run a TRT engine by specifying engine path and giving input data.
'''
class HostDeviceMem(object):
'''
Helper class to record host-device memory pointer pairs
'''
def __init__(self, host_mem, device_mem):
self.host = host_mem
self.device = device_mem
def __str__(self):
return "Host:\n" + str(self.host) + "\nDevice:\n" + str(self.device)
def __repr__(self):
return self.__str__()
def __init__(self, engine_path):
self.engine_path = engine_path
self.logger = trt.Logger(trt.Logger.WARNING)
self.runtime = trt.Runtime(self.logger)
# load and deserialize TRT engine
self.engine = self.load_engine()
# allocate input/output memory buffers
self.inputs, self.outputs, self.bindings, self.stream = self.allocate_buffers(self.engine)
# create context
self.context = self.engine.create_execution_context()
# Dict of NumPy dtype -> torch dtype (when the correspondence exists). From: https://github.com/pytorch/pytorch/blob/e180ca652f8a38c479a3eff1080efe69cbc11621/torch/testing/_internal/common_utils.py#L349
self.numpy_to_torch_dtype_dict = {
bool : torch.bool,
np.uint8 : torch.uint8,
np.int8 : torch.int8,
np.int16 : torch.int16,
np.int32 : torch.int32,
np.int64 : torch.int64,
np.float16 : torch.float16,
np.float32 : torch.float32,
np.float64 : torch.float64,
np.complex64 : torch.complex64,
np.complex128 : torch.complex128
}
def load_engine(self):
with open(self.engine_path, 'rb') as f:
engine = self.runtime.deserialize_cuda_engine(f.read())
return engine
def allocate_buffers(self, engine):
'''
Allocates all buffers required for an engine, i.e. host/device inputs/outputs.
'''
inputs = []
outputs = []
bindings = []
stream = CudaStreamContext()
for i in range(engine.num_io_tensors):
tensor_name = engine.get_tensor_name(i)
size = trt.volume(engine.get_tensor_shape(tensor_name))
dtype = trt.nptype(engine.get_tensor_dtype(tensor_name))
# Allocate host and device buffers
host_mem = np.empty(size, dtype)
device_mem = cuda_call(cudart.cudaMalloc(host_mem.nbytes))
# Append the device buffer address to device bindings. When cast to int, it's a linear index into the context's memory (like memory address).
bindings.append(int(device_mem))
# Append to the appropriate input/output list.
if engine.get_tensor_mode(tensor_name) == trt.TensorIOMode.INPUT:
inputs.append(self.HostDeviceMem(host_mem, device_mem))
else:
outputs.append(self.HostDeviceMem(host_mem, device_mem))
return inputs, outputs, bindings, stream
def __call__(self, model_inputs: list, timing=False):
'''
Inference step (like forward() in PyTorch).
model_inputs: list of numpy array or list of torch.Tensor (on GPU)
'''
NUMPY = False
TORCH = False
if isinstance(model_inputs[0], np.ndarray):
NUMPY = True
elif torch.is_tensor(model_inputs[0]):
TORCH = True
else:
assert False, 'Unsupported input data format!'
# batch size consistency check
if NUMPY:
batch_size = np.unique(np.array([i.shape[0] for i in model_inputs]))
elif TORCH:
batch_size = np.unique(np.array([i.size(dim=0) for i in model_inputs]))
assert len(batch_size) == 1, 'Input batch sizes are not consistent!'
batch_size = batch_size[0]
for i, model_input in enumerate(model_inputs):
binding_name = self.engine.get_tensor_name(i) # i-th input/output name
binding_dtype = trt.nptype(self.engine.get_tensor_dtype(binding_name)) # trt can only tell to numpy dtype
# input type cast
if NUMPY:
model_input = model_input.astype(binding_dtype)
elif TORCH:
model_input = model_input.to(self.numpy_to_torch_dtype_dict[binding_dtype])
if NUMPY:
# fill host memory with flattened input data
np.copyto(self.inputs[i].host, model_input.ravel())
elif TORCH:
nbytes = model_input.element_size() * model_input.nelement()
if timing:
memcpy_device_to_device(self.inputs[i].device, model_input.data_ptr(), nbytes)
else:
# for Torch GPU tensor it's easier, can just do Device to Device copy
memcpy_device_to_device_async(self.inputs[i].device, model_input.data_ptr(), nbytes, self.stream.stream)
if NUMPY:
if timing:
[memcpy_host_to_device(inp.device, inp.host) for inp in self.inputs]
else:
# input, Host to Device
[memcpy_host_to_device_async(inp.device, inp.host, self.stream.stream) for inp in self.inputs]
for i in range(self.engine.num_io_tensors):
self.context.set_tensor_address(self.engine.get_tensor_name(i), self.bindings[i])
duration = 0
if timing:
start_time = time()
self.context.execute_v2(bindings=self.bindings)
end_time = time()
duration = end_time - start_time
else:
# run inference
self.context.execute_async_v3(stream_handle=self.stream.stream)
if timing:
[cuda_call(cudart.cudaMemcpy(out.host.ctypes.data, out.device, out.host.nbytes, cudart.cudaMemcpyKind.cudaMemcpyDeviceToHost)) for out in self.outputs]
else:
# output, Device to Host
[memcpy_device_to_host_async(out.host, out.device, self.stream.stream) for out in self.outputs]
if not timing:
# synchronize to ensure completion of async calls
self.stream.synchronize()
if NUMPY:
return [out.host.reshape(batch_size,-1) for out in self.outputs], duration
elif TORCH:
return [torch.from_numpy(out.host.reshape(batch_size,-1)) for out in self.outputs], duration
def build_engine():
TRT_LOGGER = trt.Logger(trt.Logger.INFO)
TRT_BUILDER = trt.Builder(TRT_LOGGER)
for precision in BUILD:
engine_filename = '_'.join([MODEL_NAME, gpu_name, precision]) + '.engine'
if os.path.exists(engine_filename):
print(f'Engine file {engine_filename} exists. Skip building...')
continue
print(f'Building {precision} engine of {MODEL_NAME} model on {gpu_name} GPU...')
## parse ONNX model
network_creation_flag = 0
if "EXPLICIT_BATCH" in trt.NetworkDefinitionCreationFlag.__members__.keys():
network_creation_flag = 1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH)
network = TRT_BUILDER.create_network(network_creation_flag)
onnx_parser = trt.OnnxParser(network, TRT_LOGGER)
parse_success = onnx_parser.parse_from_file(ONNX_MODEL)
for idx in range(onnx_parser.num_errors):
print(onnx_parser.get_error(idx))
if not parse_success:
sys.exit('ONNX model parsing failed')
## build TRT engine (configuration options at: https://docs.nvidia.com/deeplearning/tensorrt/api/python_api/infer/Core/BuilderConfig.html#ibuilderconfig)
config = TRT_BUILDER.create_builder_config()
seq_len = network.get_input(0).shape[1]
# handle dynamic shape (min/opt/max): https://docs.nvidia.com/deeplearning/tensorrt/developer-guide/index.html#work_dynamic_shapes
# by default batch dim set as 1 for all min/opt/max. If there are batch need, change the value for opt and max accordingly
profile = TRT_BUILDER.create_optimization_profile()
profile.set_shape("input_ids", (1,seq_len), (1,seq_len), (1,seq_len))
profile.set_shape("attention_mask", (1,seq_len), (1,seq_len), (1,seq_len))
config.add_optimization_profile(profile)
config.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, 4096 * (1 << 20)) # 4096 MiB
# precision
if precision == 'fp32':
config.clear_flag(trt.BuilderFlag.TF32) # TF32 enabled by default, need to clear flag
elif precision == 'tf32':
pass
elif precision == 'fp16':
config.set_flag(trt.BuilderFlag.FP16)
# build
serialized_engine = TRT_BUILDER.build_serialized_network(network, config)
## save TRT engine
with open(engine_filename, 'wb') as f:
f.write(serialized_engine)
print(f'Engine is saved to {engine_filename}')
def test_engine():
for precision in TEST:
## load and deserialize TRT engine
engine_filename = '_'.join([MODEL_NAME, gpu_name, precision]) + '.engine'
print(f'Running inference on engine {engine_filename}')
model = TRTModel(engine_filename)
## psuedo-random input test
batch_size = 1
seq_len = model.engine.get_tensor_shape(model.engine.get_tensor_name(0))[1]
vocab = 128203
gpu = torch.device('cuda')
torch.manual_seed(0) # make sure in each test the seed are the same
input_ids = torch.randint(0, vocab, (batch_size, seq_len), dtype=torch.long, device=gpu)
attention_mask = torch.randint(0, 2, (batch_size, seq_len), dtype=torch.long, device=gpu)
inputs = [input_ids, attention_mask]
outputs, duration = model(inputs, timing=True)
nreps = 100
duration_total = 0
for _ in range(nreps):
outputs, duration = model(inputs, timing=True)
duration_total += duration
print(f'Average Inference time (ms) of {nreps} runs: {duration_total/nreps*1000:.3f}')
def correctness_check_engines():
for precision in CORRECTNESS:
## load and deserialize TRT engine
engine_filename1 = '_'.join([ONNX_MODEL, 'correctness_check_original', gpu_name, precision]) + '.engine'
engine_filename2 = '_'.join([ONNX_MODEL, 'correctness_check_plugin', gpu_name, precision]) + '.engine'
assert os.path.exists(engine_filename1), f'Engine file {engine_filename1} does not exist. Please build the engine first by --build'
assert os.path.exists(engine_filename2), f'Engine file {engine_filename2} does not exist. Please build the engine first by --build'
print(f'Running inference on original engine {engine_filename1} and plugin engine {engine_filename2}')
model1 = TRTModel(engine_filename1)
model2 = TRTModel(engine_filename2)
## psuedo-random input test
batch_size = 1
seq_len = model1.engine.get_tensor_shape(model1.engine.get_tensor_name(0))[1]
vocab = 128203
gpu = torch.device('cuda')
# torch.manual_seed(0) # make sure in each test the seed are the same
input_ids = torch.randint(0, vocab, (batch_size, seq_len), dtype=torch.long, device=gpu)
attention_mask = torch.randint(0, 2, (batch_size, seq_len), dtype=torch.long, device=gpu)
inputs = [input_ids, attention_mask]
outputs1, _ = model1(inputs)
outputs2, _ = model2(inputs)
# element-wise and layer-wise output comparison
for i in range(len(outputs1)):
avg_abs_error = torch.sum(torch.abs(torch.sub(outputs1[i], outputs2[i]))) / torch.numel(outputs1[i])
max_abs_error = torch.max(torch.abs(torch.sub(outputs1[i], outputs2[i])))
print(f"[Layer {i} Element-wise Check] Avgerage absolute error: {avg_abs_error.item():e}, Maximum absolute error: {max_abs_error.item():e}. 1e-2~1e-3 expected for FP16 (10 significance bits) and 1e-6~1e-7 expected for FP32 (23 significance bits)" ) # machine epsilon for different precisions: https://en.wikipedia.org/wiki/Machine_epsilon
if BUILD:
build_engine()
if TEST:
test_engine()
if CORRECTNESS:
correctness_check_engines()
demo/DeBERTa/requirements.txt
+24
View File
@@ -0,0 +1,24 @@
#
# 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.
#
-f https://download.pytorch.org/whl/cu113/torch_stable.html
torch==1.11.0+cu113
transformers==4.18.0
onnxruntime-gpu>=1.12
argparse
--extra-index-url https://pypi.ngc.nvidia.com
polygraphy
demo/Diffusion/.gitignore
+6
View File
@@ -0,0 +1,6 @@
__pycache__/
onnx/
engine/
output/
pytorch_model/
artifacts_cache/
demo/Diffusion/README.md
Executable
+538
View File
@@ -0,0 +1,538 @@
# Introduction
This demo application ("demoDiffusion") showcases the acceleration of Stable Diffusion and ControlNet pipeline using TensorRT.
# Setup
### Clone the TensorRT OSS repository
```bash
git clone git@github.com:NVIDIA/TensorRT.git -b release/11.0 --single-branch
cd TensorRT
```
### Launch NVIDIA pytorch container
Install nvidia-docker using [these intructions](https://docs.nvidia.com/datacenter/cloud-native/container-toolkit/install-guide.html#docker).
```bash
# Create a directory for persistent dependencies
mkdir -p deps
# Launch container with volume mounts
docker run --rm -it --gpus all \
-v $PWD:/workspace \
-v $PWD/deps:/workspace/deps \
nvcr.io/nvidia/pytorch:26.03-py3 /bin/bash
```
> **NOTE:** Mounting `/workspace/deps` as a volume ensures dependencies persist across container restarts. After initial installation, subsequent container launches will reuse the installed dependencies.
NOTE: The demo supports CUDA>=13.0
### Install the required packages
This demo uses a family-based dependency management system. Install dependencies for the model families you want to use:
**Install all dependencies (recommended for first-time users):**
```bash
python3 setup.py all
```
**Or install specific model families:**
```bash
# SD family: SD 1.4, SDXL, SD3, SD3.5, SVD (Stable Video Diffusion), Stable Cascade
python3 setup.py sd
# Flux family: Flux.1-dev, Flux.1-schnell, Flux.1-Canny, Flux.1-Depth, Flux.1-Kontext
python3 setup.py flux
# Cosmos family: Cosmos-Predict2 text2image, video2world
python3 setup.py cosmos
```
**Additional options:**
```bash
# Force reinstall even if already installed
python3 setup.py all --force
# Install dependencies to a custom location
# Option 1 (recommended): set the env var and install
export TENSORRT_DIFFUSION_DEPS_ROOT=/custom/path/deps
python3 setup.py all
# Option 2: install to a path without changing this shell
# Remember to export the env var in the environment that runs the demos,
# so deps.configure() can find the custom path.
python3 setup.py all --deps-root /custom/path/deps
# Then, before running any demo scripts:
export TENSORRT_DIFFUSION_DEPS_ROOT=/custom/path/deps
```
**Check installation status:**
```bash
python3 -c "from demo_diffusion import deps; deps.print_status()"
```
Check your installed TensorRT version using:
```bash
python3 -c 'import tensorrt; print(tensorrt.__version__)'
```
> NOTE: Alternatively, you can download and install TensorRT packages from [NVIDIA TensorRT Developer Zone](https://developer.nvidia.com/tensorrt).
> NOTE: demoDiffusion has been tested on systems with NVIDIA H100, A100, L40, T4, and RTX4090 GPUs, and the following software configuration.
# Running demoDiffusion
### Review usage instructions for the supported pipelines
```bash
python3 demo_txt2img.py --help
python3 demo_img2img.py --help
python3 demo_controlnet.py --help
python3 demo_txt2img_xl.py --help
python3 demo_txt2img_flux.py --help
python3 demo_txt2vid_wan.py --help
```
### HuggingFace user access token
To download model checkpoints for the Stable Diffusion pipelines, obtain a `read` access token to HuggingFace Hub. See [instructions](https://huggingface.co/docs/hub/security-tokens).
```bash
export HF_TOKEN=<your access token>
```
### Generate an image guided by a text prompt
```bash
python3 demo_txt2img.py "a beautiful photograph of Mt. Fuji during cherry blossom" --hf-token=$HF_TOKEN
```
### Faster Text-to-image using SD1.4 INT8 & FP8 quantization using ModelOpt
Run the below command to generate an image with SD1.4 in INT8
```bash
python3 demo_txt2img.py "a beautiful photograph of Mt. Fuji during cherry blossom" --hf-token=$HF_TOKEN --int8
```
Run the below command to generate an image with SD1.4 in FP8. (FP8 is only supported on Hopper and Ada.)
```bash
python3 demo_txt2img.py "a beautiful photograph of Mt. Fuji during cherry blossom" --hf-token=$HF_TOKEN --fp8
```
### Generate an image guided by an initial image and a text prompt
```bash
wget https://raw.githubusercontent.com/CompVis/stable-diffusion/main/assets/stable-samples/img2img/sketch-mountains-input.jpg -O sketch-mountains-input.jpg
python3 demo_img2img.py "A fantasy landscape, trending on artstation" --hf-token=$HF_TOKEN --input-image=sketch-mountains-input.jpg
```
### Generate an image with ControlNet guided by image(s) and text prompt(s)
```bash
python3 demo_controlnet.py "Stormtrooper's lecture in beautiful lecture hall" --controlnet-type depth --hf-token=$HF_TOKEN --denoising-steps 20 --onnx-dir=onnx-cnet-depth --engine-dir=engine-cnet-depth
```
> NOTE: `--input-image` must be a pre-processed image corresponding to `--controlnet-type`. If unspecified, a sample image will be downloaded. Supported controlnet types include: `canny`, `depth`, `hed`, `mlsd`, `normal`, `openpose`, `scribble`, and `seg`.
Examples:
<img src="https://drive.google.com/uc?export=view&id=17ub3MVSQHp26ty-wioNX6iQQ-nAveYSV" alt= “” width="800" height="400">
#### Combining multiple conditionings
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
> **NOTE:** SDXL and later Stable Diffusion models require sd dependencies to be installed. Install with: `python3 setup.py sd`
Run the below command to generate an image with Stable Diffusion XL
```bash
python3 demo_txt2img_xl.py "a photo of an astronaut riding a horse on mars" --hf-token=$HF_TOKEN --version=xl-1.0
```
The optional refiner model may be enabled by specifying `--enable-refiner` and separate directories for storing refiner onnx and engine files using `--onnx-refiner-dir` and `--engine-refiner-dir` respectively.
```bash
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.
> NOTE: FP8 quantization (`--fp8`) is supported.
### Generate an image guided by a text prompt, and using specified LoRA model weight updates
```bash
# FP16
python3 demo_txt2img_xl.py "Picture of a rustic Italian village with Olive trees and mountains" --version=xl-1.0 --lora-path "ostris/crayon_style_lora_sdxl" "ostris/watercolor_style_lora_sdxl" --lora-weight 0.3 0.7 --onnx-dir onnx-sdxl-lora --engine-dir engine-sdxl-lora --build-enable-refit
# FP8
python3 demo_txt2img_xl.py "Picture of a rustic Italian village with Olive trees and mountains" --version=xl-1.0 --lora-path "ostris/crayon_style_lora_sdxl" "ostris/watercolor_style_lora_sdxl" --lora-weight 0.3 0.7 --onnx-dir onnx-sdxl-lora --engine-dir engine-sdxl-lora --fp8
```
### Faster Text-to-image using SDXL INT8 & FP8 quantization using ModelOpt
Run the below command to generate an image with Stable Diffusion XL in INT8
```bash
python3 demo_txt2img_xl.py "a photo of an astronaut riding a horse on mars" --version xl-1.0 --onnx-dir onnx-sdxl --engine-dir engine-sdxl --int8
```
Run the below command to generate an image with Stable Diffusion XL in FP8. (FP8 is only supported on Hopper and Ada.)
```bash
python3 demo_txt2img_xl.py "a photo of an astronaut riding a horse on mars" --version xl-1.0 --onnx-dir onnx-sdxl --engine-dir engine-sdxl --fp8
```
> Note that INT8 & FP8 quantization is only supported for SDXL, and won't work with LoRA weights. FP8 quantization is only supported on Hopper and Ada. Some prompts may produce better inputs with fewer denoising steps (e.g. `--denoising-steps 20`) but this will repeat the calibration, ONNX export, and engine building processes for the U-Net.
For step-by-step tutorials to run INT8 & FP8 inference on stable diffusion models, please refer to examples in [TensorRT ModelOpt diffusers sample](https://github.com/NVIDIA/TensorRT-Model-Optimizer/tree/main/diffusers).
### Faster Text-to-Image using SDXL Turbo
Produce 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
```
### Generate an image guided by a text prompt using Stable Diffusion 3 and its variants
Run the command below to generate an image using Stable Diffusion 3 and Stable Diffusion 3.5
```bash
# Stable Diffusion 3
python3 demo_txt2img_sd3.py "A vibrant street wall covered in colorful graffiti, the centerpiece spells \"SD3 MEDIUM\", in a storm of colors" --version sd3 --hf-token=$HF_TOKEN
# Stable Diffusion 3.5-medium
python3 demo_txt2img_sd35.py "a beautiful photograph of Mt. Fuji during cherry blossom" --version=3.5-medium --denoising-steps=30 --guidance-scale 3.5 --hf-token=$HF_TOKEN --bf16 --download-onnx-models
# Stable Diffusion 3.5-large
python3 demo_txt2img_sd35.py "a beautiful photograph of Mt. Fuji during cherry blossom" --version=3.5-large --denoising-steps=30 --guidance-scale 3.5 --hf-token=$HF_TOKEN --bf16 --download-onnx-models
# Stable Diffusion 3.5-large FP8
python3 demo_txt2img_sd35.py "a beautiful photograph of Mt. Fuji during cherry blossom" --version=3.5-large --denoising-steps=30 --guidance-scale 3.5 --hf-token=$HF_TOKEN --fp8 --download-onnx-models --onnx-dir onnx_35_fp8/ --engine-dir engine_35_fp8/
```
You can also specify an input image conditioning as shown below
```bash
wget https://raw.githubusercontent.com/CompVis/latent-diffusion/main/data/inpainting_examples/overture-creations-5sI6fQgYIuo.png -O dog-on-bench.png
# Stable Diffusion 3
python3 demo_txt2img_sd3.py "dog wearing a sweater and a blue collar" --version sd3 --input-image dog-on-bench.png --hf-token=$HF_TOKEN
```
Note that a denosing-percentage is applied to the number of denoising-steps when an input image conditioning is provided. Its default value is set to 0.6. This parameter can be updated using `--denoising-percentage`
### Generate an image with Stable Diffusion v3.5-large with ControlNet guided by an image and a text prompt
```bash
# Depth BF16
python3 demo_controlnet_sd35.py "a photo of a man" --controlnet-type depth --hf-token=$HF_TOKEN --denoising-steps 40 --guidance-scale 4.5 --bf16 --download-onnx-models --low-vram
# Depth FP8
python3 demo_controlnet_sd35.py "a photo of a man" --version=3.5-large --fp8 --controlnet-type depth --download-onnx-models --denoising-steps=40 --guidance-scale 4.5 --hf-token=$HF_TOKEN --low-vram
# Canny BF16
python3 demo_controlnet_sd35.py "A Night time photo taken by Leica M11, portrait of a Japanese woman in a kimono, looking at the camera, Cherry blossoms" --controlnet-type canny --hf-token=$HF_TOKEN --denoising-steps 60 --guidance-scale 3.5 --bf16 --download-onnx-models --low-vram
# Canny FP8
python3 demo_controlnet_sd35.py "A Night time photo taken by Leica M11, portrait of a Japanese woman in a kimono, looking at the camera, Cherry blossoms" --version=3.5-large --fp8 --controlnet-type canny --hf-token=$HF_TOKEN --denoising-steps 60 --guidance-scale 3.5 --low-vram --download-onnx-models
# Blur
python3 demo_controlnet_sd35.py "generated ai art, a tiny, lost rubber ducky in an action shot close-up, surfing the humongous waves, inside the tube, in the style of Kelly Slater" --controlnet-type blur --hf-token=$HF_TOKEN --denoising-steps 60 --guidance-scale 3.5 --bf16 --download-onnx-models --low-vram
```
### Generate a video guided by an initial image using Stable Video Diffusion
Download the pre-exported ONNX model
```bash
pip install -U "huggingface_hub[cli]"
hf download stabilityai/stable-video-diffusion-img2vid-xt-1-1-tensorrt --local-dir onnx-svd-xt-1-1
```
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
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 --fp8
```
> NOTE: There is a bug in HuggingFace, you can workaround with following this [PR](https://github.com/huggingface/diffusers/pull/6562/files)
```
if torch.is_tensor(num_frames):
num_frames = num_frames.item()
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
```
NOTE: The min and max guidance scales are configured using --min-guidance-scale and --max-guidance-scale respectively.
### 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
```
> NOTE: The pipeline is only enabled for the BF16 model weights
> NOTE: The pipeline only supports ONNX export using Opset 16.
> NOTE: The denoising steps and guidance scale for the Prior and Decoder models are configured using --prior-denoising-steps, --prior-guidance-scale, --decoder-denoising-steps, and --decoder-guidance-scale respectively.
### Generating Images with Flux
> **NOTE:** Flux models require Flux family dependencies. Install with: `python3 setup.py flux`
#### 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 --download-onnx-models
# FP8
python3 demo_txt2img_flux.py "a beautiful photograph of Mt. Fuji during cherry blossom" --hf-token=$HF_TOKEN --quantization-level 4 --fp8 --download-onnx-models
# 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 --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" --quantization-level 4 --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 --download-onnx-models
```
---
#### 2. Generate an Image from an Initial Image + Text Prompt
Download an example input image:
```bash
wget "https://miro.medium.com/v2/resize:fit:640/format:webp/1*iD8mUonHMgnlP0qrSx3qPg.png" -O yellow.png
```
Run the image-to-image pipeline:
```bash
python3 demo_img2img_flux.py "A home with 2 floors and windows. The front door is purple" --hf-token=$HF_TOKEN --input-image yellow.png --image-strength 0.95 --bf16 --onnx-dir onnx-flux-dev/bf16 --engine-dir engine-flux-dev/
```
---
#### 3. Generate an Image Using Flux ControlNet
##### Download the Control Image
```bash
wget https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/robot.png
```
##### Calibration Data for native ONNX export (FP8 Pipeline)
FP8 ControlNet pipelines require downloading a calibration dataset and providing the path. You can use the datasets provided by Black Forest Labs here: [depth](https://drive.google.com/file/d/1DFfhOSrTlKfvBFLcD2vAALwwH4jSGdGk/view) | [canny](https://drive.google.com/file/d/1dRoxOL-vy3tSAesyqBSJoUWsbkMwv3en/view)
You can use the `--calibraton-dataset` flag to specify the path, which is set to `./{depth/canny}-eval/benchmark` by default if not provided. Note that the dataset should have `inputs/` and `prompts/` underneath the provided path, matching the format of the BFL dataset.
##### Depth ControlNet
```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 --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 --download-onnx-models --build-static-batch --quantization-level 4
# 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 --quantization-level 4 --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 --download-onnx-models --build-static-batch
```
##### 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 --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 --denoising-steps 30 --download-onnx-models --build-static-batch --quantization-level 4
# 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 --quantization-level 4 --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 --denoising-steps 30 --download-onnx-models --build-static-batch
```
#### 4. Generate an Image Using Flux LoRA
FLUX supports loading LoRA for Flux.1-Dev and Flux.1-Schnell. Make sure the target lora is compatible with the transformer model. Below is an example of using a [water color Flux LoRA](https://huggingface.co/SebastianBodza/flux_lora_aquarel_watercolor)
```bash
# FP16
python3 demo_txt2img_flux.py "A painting of a barista creating an intricate latte art design, with the 'Coffee Creations' logo skillfully formed within the latte foam. In a watercolor style, AQUACOLTOK. White background." --hf-token=$HF_TOKEN --lora-path "SebastianBodza/flux_lora_aquarel_watercolor" --lora-weight 1.0 --onnx-dir=onnx-flux-lora --engine-dir=engine-flux-lora
# FP8
python3 demo_txt2img_flux.py "A painting of a barista creating an intricate latte art design, with the 'Coffee Creations' logo skillfully formed within the latte foam. In a watercolor style, AQUACOLTOK. White background." --hf-token=$HF_TOKEN --lora-path "SebastianBodza/flux_lora_aquarel_watercolor" --lora-weight 1.0 --onnx-dir=onnx-flux-lora --engine-dir=engine-flux-lora --fp8
```
#### 5. Edit an Image using Flux Kontext
```bash
wget https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/cat.png
# BF16
python3 demo_img2img_flux.py "Add a hat to the cat" --version="flux.1-kontext-dev" --hf-token=$HF_TOKEN --guidance-scale 2.5 --kontext-image cat.png --denoising-steps 28 --bf16 --onnx-dir onnx-kontext --engine-dir engine-kontext --download-onnx-models
# FP8
python3 demo_img2img_flux.py "Add a hat to the cat" --version="flux.1-kontext-dev" --hf-token=$HF_TOKEN --guidance-scale 2.5 --kontext-image cat.png --denoising-steps 28 --fp8 --onnx-dir onnx-kontext-fp8 --engine-dir engine-kontext-fp8 --download-onnx-models --quantization-level 4
# FP4
python3 demo_img2img_flux.py "Add a hat to the cat" --version="flux.1-kontext-dev" --hf-token=$HF_TOKEN --guidance-scale 2.5 --kontext-image cat.png --denoising-steps 28 --fp4 --onnx-dir onnx-kontext-fp4 --engine-dir engine-kontext-fp4 --download-onnx-models
```
---
#### 5. Export ONNX Models Only (Skip Inference)
Use the `--onnx-export-only` flag to export ONNX models on a higher-VRAM device. The exported ONNX models can be used on a device with lower VRAM for the engine build and inference steps.
```bash
python3 demo_txt2img_flux.py "a beautiful photograph of Mt. Fuji during cherry blossom" --hf-token=$HF_TOKEN --onnx-export-only
```
---
#### 6. Running Flux on GPUs with Limited Memory
##### Optimization Flags
- `--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 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
Memory usage captured below excludes the ONNX export step, and assumes use of the `--build-static-batch` flag to reduce activation VRAM usage. Users can either use [pre-exported ONNX models](README.md#download-pre-exported-models-recommended-for-48gb-vram) or export the models separately on a higher-VRAM device using [--onnx-export-only](README.md#4-export-onnx-models-only-skip-inference).
| Precision | Default VRAM Usage | With `--low-vram` |
| --------- | ------------------ | ----------------- |
| FP16 | 39.3 GB | 23.9 GB |
| BF16 | 35.7 GB | 23.9 GB |
| FP8 | 24.6 GB | 14.9 GB |
| FP4 | 21.67 GB | 11.1 GB |
NOTE: The FP8 and FP4 Pipelines are supported on Hopper/Ada/Blackwell devices only. The FP4 pipeline is most performant on Blackwell devices.
### Run Cosmos2 World Foundation Models
> **NOTE:** Cosmos models require Cosmos family dependencies. Install with: `python3 setup.py cosmos`
Select the prompts and export them as below
```bash
export PROMPT="A close-up shot captures a vibrant yellow scrubber vigorously working on a grimy plate, its bristles moving in circular motions to lift stubborn grease and food residue. The dish, once covered in remnants of a hearty meal, gradually reveals its original glossy surface. Suds form and bubble around the scrubber, creating a satisfying visual of cleanliness in progress. The sound of scrubbing fills the air, accompanied by the gentle clinking of the dish against the sink. As the scrubber continues its task, the dish transforms, gleaming under the bright kitchen lights, symbolizing the triumph of cleanliness over mess."
export NEGATIVE_PROMPT="The video captures a series of frames showing ugly scenes, static with no motion, motion blur, over-saturation, shaky footage, low resolution, grainy texture, pixelated images, poorly lit areas, underexposed and overexposed scenes, poor color balance, washed out colors, choppy sequences, jerky movements, low frame rate, artifacting, color banding, unnatural transitions, outdated special effects, fake elements, unconvincing visuals, poorly edited content, jump cuts, visual noise, and flickering. Overall, the video is of poor quality."
```
#### 1. Generate an Image from a Text Prompt
##### Run Cosmos-Predict2-2B-Text2Image
```bash
# BF16
python3 demo_txt2image_cosmos.py "$PROMPT" --negative-prompt="$NEGATIVE_PROMPT" --hf-token=$HF_TOKEN
```
#### 2. Generate a Video guided by an Initial Video Conditioning and a Text Prompt
##### Run Cosmos-Predict2-2B-Video2World (only PyTorch backend enabled)
```bash
# BF16
python3 demo_vid2world_cosmos.py "$PROMPT" --negative-prompt="$NEGATIVE_PROMPT" --hf-token=$HF_TOKEN
```
### Specify Custom Paths for ONNX models and TensorRT engines (FLUX, Stable Diffusion 3.5 and Cosmos only)
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`.
### Generate a video from a text prompt using Wan
Run the below command to generate 81 frames of video at 720×1280 resolution using Wan 2.2. Due to the high memory requirements of this model, it is recommended to enable `--low-vram` and use a Blackwell device.
```bash
# Default (81 frames, 720x1280) with --low-vram enabled
python3 demo_txt2vid_wan.py "A serene bamboo forest with sunlight filtering through the leaves" --hf-token=$HF_TOKEN --low-vram
# Adjust denoising steps, guidance scales, negative prompt, seed, warmup runs
python3 demo_txt2vid_wan.py "Ocean waves crashing on a beach at sunset" --hf-token=$HF_TOKEN --low-vram --denoising-steps 50 --guidance-scale 4.5 --guidance-scale-2 3.5 --negative-prompt "blurry, low quality, static" --seed 42 --num-warmup-runs 0
```
## 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_data/calibration-images/rocket.png
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 482 KiB

demo/Diffusion/calibration_data/calibration-prompts.txt
+1079
View File
File diff suppressed because it is too large Load Diff
demo/Diffusion/demo_controlnet.py
+161
View File
@@ -0,0 +1,161 @@
#
# 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.
#
# Configure dependencies before any external imports
from demo_diffusion import deps
deps.configure("sd")
import argparse
import controlnet_aux
import torch
from cuda.bindings import runtime as cudart
from PIL import 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 ControlNet Demo", conflict_handler='resolve')
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']")
parser.add_argument('--controlnet-scale', nargs='+', type=float, default=[1.0], help="The outputs of the controlnet are multiplied by `controlnet_scale` before they are added to the residual in the original unet, can be `None`, `float` or `float` list")
return parser.parse_args()
if __name__ == "__main__":
print("[I] Initializing StableDiffusion controlnet demo using TensorRT")
args = parseArgs()
# 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)
# Check images
input_images = []
if len(args.input_image) > 0:
for image in args.input_image:
input_images.append(Image.open(image))
else:
for controlnet in args.controlnet_type:
if controlnet == "canny":
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.4 and SDXL base pipelines only. Version provided: {args.version}"
)
input_images.append(canny_image.resize((args.width, args.height)))
elif controlnet == "normal":
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.width, args.height)))
elif controlnet == "depth":
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.width, args.height)))
elif controlnet == "hed":
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.width, args.height)))
elif controlnet == "mlsd":
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.width, args.height)))
elif controlnet == "openpose":
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.width, args.height)))
elif controlnet == "scribble":
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.width, args.height)))
elif controlnet == "seg":
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.width, args.height)))
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 = dd_argparse.process_pipeline_args(args)
# Initialize demo
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,
)
# Load TensorRT engines and pytorch modules
demo.loadEngines(
args.engine_dir,
args.framework_model_dir,
args.onnx_dir,
**kwargs_load_engine)
# Load resources
_, shared_device_memory = cudart.cudaMalloc(demo.calculateMaxDeviceMemory())
demo.activateEngines(shared_device_memory)
demo.loadResources(args.height, args.width, args.batch_size, args.seed)
# Run inference
demo_kwargs = {'input_image': input_images, 'controlnet_scales': controlnet_scale}
demo.run(*args_run_demo, **demo_kwargs)
demo.teardown()
demo/Diffusion/demo_controlnet_sd35.py
+191
View File
@@ -0,0 +1,191 @@
#
# 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.
#
# Configure dependencies before any external imports
from demo_diffusion import deps
deps.configure("sd")
import argparse
import torch
from cuda.bindings import runtime as cudart
from PIL import 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 3.5-large ControlNet Demo", conflict_handler="resolve"
)
parser = dd_argparse.add_arguments(parser)
parser.add_argument(
"--version",
type=str,
default="3.5-large",
choices={"3.5-large"},
help="Version of Stable Diffusion 3.5",
)
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)")
parser.add_argument(
"--max-sequence-length",
type=int,
default=256,
help="Maximum sequence length to use with the prompt.",
)
parser.add_argument(
"--control-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",
type=str,
default="canny",
help="Controlnet type (single type only), can be 'canny', 'depth', 'blur', etc.",
)
parser.add_argument(
"--controlnet-scale",
type=float,
default=1.0,
help="The outputs of the controlnet are multiplied by `controlnet_scale` before they are added to the residual in the original Transformer",
)
return parser.parse_args()
def process_demo_args(args):
batch_size = args.batch_size
prompt = args.prompt
negative_prompt = args.negative_prompt
# Process prompt
if not isinstance(prompt, list):
raise ValueError(f"`prompt` must be of type `str` list, but is {type(prompt)}")
prompt = prompt * batch_size
if not isinstance(negative_prompt, list):
raise ValueError(f"`--negative-prompt` must be of type `str` list, but is {type(negative_prompt)}")
if len(negative_prompt) == 1:
negative_prompt = negative_prompt * batch_size
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`"
)
# Controlnet configuration
if not isinstance(args.controlnet_type, str):
raise ValueError(f"`--controlnet-type` must be of type `str`, but is {type(args.controlnet_type)}")
# Controlnet configuration
if not isinstance(args.controlnet_scale, float):
raise ValueError(f"`--controlnet-scale` must be of type `float`, but is {type(args.controlnet_scale)}")
# Convert controlnet scales to tensor
controlnet_scale = torch.tensor(args.controlnet_scale)
# Check images
input_images = []
if len(args.control_image) > 0:
for image in args.control_image:
input_images.append(Image.open(image))
else:
if args.controlnet_type == "canny":
canny_image = image_module.download_image(
"https://huggingface.co/datasets/diffusers/diffusers-images-docs/resolve/main/canny.png"
)
input_images.append(canny_image.resize((args.width, args.height)))
elif args.controlnet_type == "depth":
depth_image = image_module.download_image(
"https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/marigold/marigold_einstein_lcm_depth.png"
)
input_images.append(depth_image.resize((args.width, args.height)))
elif args.controlnet_type == "blur":
blur_image = image_module.download_image(
"https://huggingface.co/datasets/diffusers/diffusers-images-docs/resolve/main/blur.png"
)
input_images.append(blur_image.resize((args.width, args.height)))
else:
raise ValueError(f"You should implement the conditonal image of this controlnet: {args.controlnet_type}")
assert len(input_images) > 0
kwargs_run_demo = {
"prompt": prompt,
"negative_prompt": negative_prompt,
"height": args.height,
"width": args.width,
"control_image": input_images,
"controlnet_scale": controlnet_scale,
"batch_count": args.batch_count,
"num_warmup_runs": args.num_warmup_runs,
"use_cuda_graph": args.use_cuda_graph,
}
return kwargs_run_demo
if __name__ == "__main__":
print("[I] Initializing StableDiffusion ControlNet demo using TensorRT")
args = parseArgs()
# Initialize demo
_, kwargs_load_engine, _ = dd_argparse.process_pipeline_args(args)
kwargs_run_demo = process_demo_args(args)
# Initialize demo
demo = pipeline_module.StableDiffusion35Pipeline.FromArgs(
args,
pipeline_type=pipeline_module.PIPELINE_TYPE.CONTROLNET,
)
# Load TensorRT engines and pytorch modules
demo.load_engines(
framework_model_dir=args.framework_model_dir,
**kwargs_load_engine,
)
if demo.low_vram:
demo.device_memory_sizes = demo.get_device_memory_sizes()
else:
_, shared_device_memory = cudart.cudaMalloc(demo.calculate_max_device_memory())
demo.activate_engines(shared_device_memory)
# Load resources
demo.load_resources(
image_height=args.height,
image_width=args.width,
batch_size=args.batch_size,
seed=args.seed,
)
# Run inference
demo.run(**kwargs_run_demo)
demo.teardown()
demo/Diffusion/demo_diffusion/__init__.py
View File
demo/Diffusion/demo_diffusion/dd_argparse.py
+465
View File
@@ -0,0 +1,465 @@
#
# 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.4",
choices=(
"1.4",
"dreamshaper-7",
"xl-1.0",
"xl-turbo",
"svd-xt-1.1",
"sd3",
"3.5-medium",
"3.5-large",
"cascade",
"flux.1-dev",
"flux.1-schnell",
"flux.1-dev-canny",
"flux.1-dev-depth",
"flux.1-kontext-dev",
"cosmos-predict2-2b-text2image",
"cosmos-predict2-14b-text2image",
"cosmos-predict2-2b-video2world",
"cosmos-predict2-14b-video2world",
),
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'",
)
parser.add_argument("--bf16", action="store_true", help="Run pipeline in BFloat16 precision")
# ONNX export
parser.add_argument(
"--onnx-opset",
type=int,
default=19,
choices=range(7, 24),
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, Stable Diffusion 3.5-large, Cosmos 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(
"--onnx-export-only",
action="store_true",
help="If set, only performs the export of models to ONNX, skipping engine build and inference.",
)
parser.add_argument(
"--download-onnx-models",
action="store_true",
help=("[FLUX and Stable Diffusion 3.5-large 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."
),
)
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."
)
parser.add_argument("--ws", action="store_true", help="Build TensorRT engines with weight streaming enabled.")
# 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",
)
# 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(
"--torch-fallback",
default=None,
type=str,
help="[FLUX, SD3.5, and Wan] Comma separated list of models to be inferenced using PyTorch instead of TRT. For example --torch-fallback text_encoder,transformer,transformer_2. If --torch-inference set, this parameter will be ignored.",
)
parser.add_argument(
"--low-vram",
action="store_true",
help="[FLUX, SD3.5, and Wan] Optimize for low VRAM usage, possibly at the expense of inference performance. Disabled by default.",
)
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
is_flux = args.version.startswith("flux")
is_sd35 = args.version.startswith("3.5")
is_wan = args.version.startswith("wan")
is_cosmos = args.version.startswith("cosmos")
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:
# optimization level set to 3 for all Flux pipelines to reduce GPU memory usage
if args.int8 or args.fp8 and not is_flux:
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")):
raise ValueError("int8 quantization is only supported for SDXL and SD1.4 pipelines.")
# fp8 support validation
if args.fp8:
# Check version compatibility
supported_versions = ("xl", "1.4", "3.5-large")
if not (any(args.version.startswith(prefix) for prefix in supported_versions) or is_flux):
raise ValueError(
"fp8 quantization is only supported for SDXL, SD1.4, SD3.5-large and FLUX pipelines."
)
# Check controlnet compatibility
if getattr(args, "controlnet_type", None) is not None:
if args.version not in ("xl-1.0", "3.5-large"):
raise ValueError("fp8 controlnet quantization is only supported for SDXL and SD3.5-large.")
if args.version == "3.5-large" and args.controlnet_type == "blur":
raise ValueError("Blur controlnet type is not supported for SD3.5.")
# Check for conflicting quantization
if args.int8:
raise ValueError("Cannot apply both int8 and fp8 quantization, please choose only one.")
# Check GPU compute capability
if sm_version < 89:
raise ValueError(
f"Cannot apply FP8 quantization for GPU with compute capability {sm_version / 10.0}. A minimum compute capability of 8.9 is required."
)
# Check SD3.5-large specific requirement
if args.version == "3.5-large" and not args.download_onnx_models:
raise ValueError(
"Native FP8 quantization is not supported for SD3.5-large. Please pass --download-onnx-models."
)
# 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"[W] The default quantization level has been set to {level} for {dtype_str}.")
if args.fp8:
# L4 fp8 fMHA on Hopper not yet enabled.
if sm_version == 90 and is_flux:
override_quant_level(3.0, "FP8")
else:
override_quant_level(3.0 if args.version == "1.4" else 4.0, "FP8")
elif args.int8:
override_quant_level(3.0, "INT8")
if args.version.startswith("flux") and 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 the Flux pipeline
assert is_flux, "FP4 precision is only supported for the Flux pipeline"
# Handle LoRA
# FLUX canny and depth official LoRAs are not supported because they modify the transformer architecture, conflicting with refit
if args.lora_path and not any(args.version.startswith(prefix) for prefix in ("xl", "flux.1-dev", "flux.1-schnell")):
raise ValueError("LoRA adapter support is only supported for SDXL, FLUX.1-dev and FLUX.1-schnell 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}")
# Force lora merge when fp8 or int8 is used with LoRA
if args.build_enable_refit and args.lora_path and (args.int8 or args.fp8):
raise ValueError(
"Engine refit should not be enabled for quantized models with LoRA. ModelOpt recommends fusing the LoRA to the model before quantization. \
See https://github.com/NVIDIA/TensorRT-Model-Optimizer/tree/main/examples/diffusers/quantization#lora"
)
# Torch-fallback and Torch-inference
if args.torch_fallback and not args.torch_inference:
assert (
is_flux or is_sd35 or is_wan or is_cosmos
), "PyTorch Fallback is only supported for Flux, Stable Diffusion 3.5, Wan and Cosmos pipelines."
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."
)
args.torch_fallback = None
# low-vram
if args.low_vram:
assert (
is_flux or is_sd35 or is_wan or is_cosmos
), "low-vram mode is only supported for Flux, Stable Diffusion 3.5, Wan and Cosmos pipelines."
# Disable SDXL LCM pipeline
if args.version == "xl-1.0" and args.scheduler == "LCM":
raise ValueError("SDXL pipeline does not support the LCM scheduler currently. Please use a different scheduler.")
# 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,
"onnx_export_only": args.onnx_export_only,
"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/deps.py
+270
View File
@@ -0,0 +1,270 @@
#
# 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.
#
"""Dependency management for TensorRT diffusion demos.
Adds the right group's site-packages to sys.path so imports work.
"""
import os
import sys
# Marker file to indicate successful installation
# Must match the marker file used by setup.py
INSTALL_COMPLETE_MARKER = ".install_complete"
# Descriptions for user-facing messages
GROUP_DESCRIPTIONS = {
"sd": "SD family (SD 1.4, SDXL, SD3, SD3.5, SVD, Stable Cascade)",
"flux": "Flux family (Black Forest Labs)",
"cosmos": "Cosmos family (NVIDIA), Wan2.2 T2V",
}
# Valid dependency groups
VALID_GROUPS = list(GROUP_DESCRIPTIONS.keys())
__all__ = [
"configure",
"get_configured_groups",
"print_status",
]
def _resolve_deps_root(deps_root: str | None) -> str:
"""Resolve the dependency root path with precedence: arg > env > default."""
if deps_root is not None:
return deps_root
return os.environ.get("TENSORRT_DIFFUSION_DEPS_ROOT", "/workspace/deps")
def _prepend_env_path(var: str, path: str, clean_root: str | None = None):
"""Prepend `path` to an os.pathsep-separated env var so child processes
(e.g. the `polygraphy` CLI) inherit the group's dependencies. If
`clean_root` is given, drop existing entries under it first."""
def _norm(p: str) -> str:
return os.path.abspath(os.path.expanduser(p))
existing = [p for p in os.environ.get(var, "").split(os.pathsep) if p]
if clean_root is not None:
root_abs = _norm(clean_root).rstrip(os.sep) + os.sep
existing = [p for p in existing if not _norm(p).startswith(root_abs)]
existing = [p for p in existing if _norm(p) != _norm(path)]
os.environ[var] = os.pathsep.join([path, *existing])
def _clean_diffusion_paths(deps_root: str = "/workspace/deps"):
"""Drop any existing paths under deps_root from sys.path."""
# Filter out any paths that live under deps_root (robust to path forms)
root_abs = os.path.abspath(os.path.expanduser(deps_root)).rstrip(os.sep) + os.sep
sys.path[:] = [
p for p in sys.path
if not os.path.abspath(os.path.expanduser(p)).startswith(root_abs)
]
def configure(
group: str,
deps_root: str | None = None,
verbose: bool = False,
clean: bool = True,
fallback: bool = False
):
"""
Configure sys.path to use dependencies from the specified group.
This function should be called at the top of each demo script, before
any other imports that depend on external packages.
Args:
group: Dependency group name ("sd", "flux", or "cosmos")
deps_root: Root directory where dependencies are installed.
If None, uses TENSORRT_DIFFUSION_DEPS_ROOT environment variable,
or defaults to "/workspace/deps"
verbose: Print configuration info (default: False)
clean: Remove other dependency group paths first (default: True)
This prevents sys.path inflation in long-running processes.
fallback: If True, gracefully handle missing dependencies instead of raising
RuntimeError. Useful for test environments using requirements.txt.
If False but USE_REQUIREMENTS=1 is set,
fallback will be automatically enabled.
Example:
from demo_diffusion import deps
deps.configure("flux")
# Or with custom location
deps.configure("flux", deps_root="/custom/path/deps")
# Or via environment variable
# export TENSORRT_DIFFUSION_DEPS_ROOT=/custom/path/deps
# deps.configure("flux")
# For test environments using requirements.txt (installed via setup.sh)
# export USE_REQUIREMENTS=1
# deps.configure("sd") # Will automatically use fallback mode
"""
if group not in VALID_GROUPS:
raise ValueError(
f"Invalid dependency group: '{group}'\n"
f"Valid groups: {', '.join(sorted(VALID_GROUPS))}"
)
# Auto-enable fallback mode if using requirements.txt installation
if not fallback and os.environ.get("USE_REQUIREMENTS") == "1":
fallback = True
if verbose:
print("Detected USE_REQUIREMENTS=1, enabling fallback mode")
# Resolve deps_root consistently across the module
deps_root = _resolve_deps_root(deps_root)
# Clean old dependency paths first (prevents inflation)
if clean:
_clean_diffusion_paths(deps_root)
# Determine Python version
python_version = f"{sys.version_info.major}.{sys.version_info.minor}"
# Construct path to site-packages for this group
deps_path = os.path.join(
deps_root,
group,
"lib",
f"python{python_version}",
"site-packages"
)
# Check if installation is complete
group_dir = os.path.join(deps_root, group)
marker_file = os.path.join(group_dir, INSTALL_COMPLETE_MARKER)
if os.path.exists(deps_path) and os.path.exists(marker_file):
# Insert at the beginning to override any system packages
sys.path.insert(0, deps_path)
# Mirror onto the environment so child processes (e.g. the polygraphy
# CLI that engine builds shell out to) use these deps too. When clean,
# drop other groups' entries under deps_root, matching sys.path above.
clean_root = deps_root if clean else None
_prepend_env_path("PYTHONPATH", deps_path, clean_root=clean_root)
_prepend_env_path("PATH", os.path.join(group_dir, "bin"), clean_root=clean_root)
if verbose:
description = GROUP_DESCRIPTIONS.get(group, group)
print(f"Configured dependencies: {description}")
print(f" Path: {deps_path}")
else:
# Dependencies not found or installation incomplete
description = GROUP_DESCRIPTIONS.get(group, group)
# Check if it's an incomplete installation
if os.path.exists(group_dir) and not os.path.exists(marker_file):
error_msg = (
f"Dependencies for '{group}' are incomplete!\n"
f" Location: {group_dir}\n"
f" Description: {description}\n"
f" A previous installation failed or was interrupted.\n"
f" To fix, run:\n"
f" python setup.py {group}\n"
f" (This will clean up and reinstall)"
)
if fallback:
if verbose:
print(f"Warning: {error_msg}")
print("Continuing with fallback mode...")
return
else:
raise RuntimeError(error_msg)
else:
# Not installed at all
error_msg = (
f"Dependencies for '{group}' not found!\n"
f" Expected at: {deps_path}\n"
f" Description: {description}\n"
f" To install, run: python setup.py {group}\n"
f" If you installed elsewhere, set TENSORRT_DIFFUSION_DEPS_ROOT or pass deps_root."
)
if fallback:
if verbose:
print(f"Warning: {error_msg}")
print("Continuing with fallback mode (assuming requirements.txt setup)...")
return
else:
raise RuntimeError(error_msg)
def get_configured_groups(deps_root: str | None = None) -> list[str]:
"""
Get list of dependency groups that are currently installed.
A group is considered installed only if both:
1. The group directory exists
2. The .install_complete marker file exists in that directory
Args:
deps_root: Root directory where dependencies are installed.
If None, uses TENSORRT_DIFFUSION_DEPS_ROOT environment variable,
or defaults to "/workspace/deps"
Returns:
List of installed group names (e.g., ["sd", "flux"])
"""
# Resolve deps_root
deps_root = _resolve_deps_root(deps_root)
installed = []
for group in VALID_GROUPS:
# Check if the group directory exists and has the completion marker
group_dir = os.path.join(deps_root, group)
marker_file = os.path.join(group_dir, INSTALL_COMPLETE_MARKER)
# Only consider it installed if marker file exists
if os.path.exists(group_dir) and os.path.isdir(group_dir) and os.path.exists(marker_file):
installed.append(group)
return sorted(installed)
def print_status(deps_root: str | None = None):
"""
Print the status of all dependency groups.
Args:
deps_root: Root directory where dependencies are installed.
If None, uses TENSORRT_DIFFUSION_DEPS_ROOT environment variable,
or defaults to "/workspace/deps"
"""
# Resolve deps_root
deps_root = _resolve_deps_root(deps_root)
print("Dependency Groups Status:")
print("-" * 60)
print(f"Location: {deps_root}")
print("-" * 60)
installed = get_configured_groups(deps_root)
for group in sorted(VALID_GROUPS):
description = GROUP_DESCRIPTIONS.get(group, group)
status = "Installed" if group in installed else "Not installed"
print(f" {group:8} - {status:15} - {description}")
print("-" * 60)
if not installed:
print("No dependency groups installed.")
print("Run: python setup.py all")
else:
print(f"{len(installed)} group(s) installed: {', '.join(installed)}")
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
+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 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.bindings import runtime as 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,
tf32=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}")
# Base command
build_command = [f"polygraphy convert {onnx_path} --convert-to trt --output {self.engine_path}"]
# Build arguments
build_args = [
"--strongly-typed",
"--tf32" if tf32 else "",
"--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, verbose=True):
if self.engine is not None:
if verbose:
print(f"Unloading TensorRT engine: {self.engine_path}")
del self.engine
self.engine = None
gc.collect()
else:
if verbose:
print(f"[W]: Unload an unloaded engine {self.engine_path}, skip unloading")
def activate(self, device_memory=None):
if device_memory is not None:
self.context = self.engine.create_execution_context(
trt.ExecutionContextAllocationStrategy.USER_MANAGED
)
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):
if not self.engine:
return
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
+111
View File
@@ -0,0 +1,111 @@
#
# 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.controlnet import SD3ControlNet
from demo_diffusion.model.diffusion_transformer import (
CosmosTransformerModel,
FluxTransformerModel,
SD3_MMDiTModel,
SD3TransformerModel,
WanTransformerModel,
)
from demo_diffusion.model.gan import VQGANModel
from demo_diffusion.model.load import unload_torch_model
from demo_diffusion.model.lora import FLUXLoraLoader, SDLoraLoader, 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 (
UNet2DConditionControlNetModel,
UNetCascadeModel,
UNetModel,
UNetTemporalModel,
UNetXLModel,
UNetXLModelControlNet,
)
from demo_diffusion.model.vae import (
AutoencoderKLWanEncoderModel,
AutoencoderKLWanModel,
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",
"CosmosTransformerModel",
# diffusion_transformer
"SD3_MMDiTModel",
"FluxTransformerModel",
"SD3TransformerModel",
"SD3ControlNet",
"WanTransformerModel",
# gan
"VQGANModel",
# lora
"SDLoraLoader",
"FLUXLoraLoader",
"merge_loras",
# scheduler
"make_scheduler",
# t5
"T5Model",
# tokenizer
"make_tokenizer",
# unet
"UNetModel",
"UNetXLModel",
"UNetXLModelControlNet",
"UNet2DConditionControlNetModel",
"UNetTemporalModel",
"UNetCascadeModel",
# vae
"VAEModel",
"SD3_VAEDecoderModel",
"VAEDecTemporalModel",
"TorchVAEEncoder",
"VAEEncoderModel",
"SD3_VAEEncoderModel",
"AutoencoderKLWanModel",
"AutoencoderKLWanEncoderModel",
# load
"unload_torch_model",
]
demo/Diffusion/demo_diffusion/model/base_model.py
+320
View File
@@ -0,0 +1,320 @@
#
# 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 json
import os
import numpy as np
import onnx
import torch
from diffusers import DiffusionPipeline
from onnx import numpy_helper
from demo_diffusion.model import load, optimizer
from demo_diffusion.model.lora import merge_loras
class BaseModel:
def __init__(
self,
version="1.4",
pipeline=None,
device="cuda",
hf_token="",
verbose=True,
framework_model_dir="pytorch_model",
fp16=False,
tf32=False,
bf16=False,
int8=False,
fp8=False,
fp4=False,
max_batch_size=16,
text_maxlen=77,
embedding_dim=768,
compression_factor=8,
):
self.name = self.__class__.__name__
self.pipeline_type = pipeline
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 = True
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.fp4 = fp4
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 = 1360 # max image resolution: 1360x1360
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,
opt_num_frames=None,
custom_model=None,
enable_lora_merge=False,
static_shape=False,
lora_loader=None,
dynamo=False,
):
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)
export_kwargs = {}
if dynamo:
export_kwargs["dynamic_shapes"] = self.get_dynamic_axes()
else:
export_kwargs["dynamic_axes"] = self.get_dynamic_axes()
inputs = self.get_sample_input(
1, opt_image_height, opt_image_width, static_shape,
**({'num_frames': opt_num_frames} if opt_num_frames else {})
)
with torch.no_grad():
torch.onnx.export(
model,
inputs,
onnx_path,
export_params=True,
do_constant_folding=self.do_constant_folding,
input_names=self.get_input_names(),
output_names=self.get_output_names(),
verbose=False,
dynamo=dynamo,
opset_version=onnx_opset,
**export_kwargs,
)
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
model = self.get_model()
with torch.inference_mode(), torch.autocast("cuda", enabled=do_autocast):
export_onnx(model)
del model
gc.collect()
torch.cuda.empty_cache()
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, version=self.version)
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")
elif self.bf16:
# Cast Resize I/O for strongly-typed TRT builds: BF16 -> FP32 inputs, FP32 -> BF16 outputs.
# TRT does not support BF16 for the Resize operator.
opt.infer_shapes()
opt.cast_resize_io(output_dtype=onnx.TensorProto.BFLOAT16)
opt.info(self.name + ": cast resize I/O for bf16")
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("modify_int8_graph", False):
opt.modify_int8_graph()
opt.info(self.name + ": modify int8 graph")
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, num_frames=None):
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
if num_frames:
latent_frames = (self.num_frames - 1) // self.temporal_compression_factor + 1
return (latent_height, latent_width, latent_frames)
return (latent_height, latent_width)
def get_minmax_dims(self, batch_size, image_height, image_width, static_batch, static_shape, num_frames=None):
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
frame_dims = ()
if num_frames:
latent_frames = (num_frames - 1) // self.temporal_compression_factor + 1
min_latent_frames = latent_frames if static_shape else self.min_latent_frames
max_latent_frames = latent_frames if static_shape else self.max_latent_frames
frame_dims = (min_latent_frames, max_latent_frames)
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,
*frame_dims
)
demo/Diffusion/demo_diffusion/model/clip.py
+555
View File
@@ -0,0 +1,555 @@
#
# 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, subfolder: str) -> int:
"""Return the embedding dimension of a CLIP with projection model."""
if version in ("xl-1.0", "xl-turbo", "cascade"):
return 1280
elif version in {"3.5-medium", "3.5-large"} and subfolder == "text_encoder":
return 768
elif version in {"3.5-medium", "3.5-large"} and subfolder == "text_encoder_2":
return 1280
else:
raise ValueError(f"Invalid version {version} + subfolder {subfolder}")
def get_clip_embedding_dim(version, pipeline):
if version in (
"1.4",
"dreamshaper-7",
"flux.1-dev",
"flux.1-schnell",
"flux.1-dev-canny",
"flux.1-dev-depth",
"flux.1-kontext-dev",
):
return 768
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,
attn_implementation="eager",
**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, version=self.version)
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, subfolder),
output_hidden_states=output_hidden_states,
)
self.subfolder = subfolder
def get_model(self, torch_inference=""):
model_opts = {"variant": "fp16", "torch_dtype": torch.float16} if self.fp16 else {"torch_dtype": torch.bfloat16}
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,
attn_implementation="eager",
**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, version=self.version)
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)
return model
demo/Diffusion/demo_diffusion/model/controlnet.py
+223
View File
@@ -0,0 +1,223 @@
#
# 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 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 = ["SD3Transformer2DModel", "SD3ControlNetModel"]
for model in models_to_import:
globals()[model] = import_from_diffusers(model, "diffusers.models")
class SD3ControlNetWrapper(torch.nn.Module):
def __init__(self, controlnet):
super().__init__()
self.controlnet = controlnet
def forward(self, hidden_states, controlnet_cond, conditioning_scale, pooled_projections, timestep):
params = {
"hidden_states": hidden_states,
"pooled_projections": pooled_projections,
"timestep": timestep,
"controlnet_cond": controlnet_cond,
"conditioning_scale": conditioning_scale,
}
out = self.controlnet(**params)["controlnet_block_samples"]
return torch.stack(out, dim=0)
class SD3ControlNet(base_model.BaseModel):
def __init__(
self,
version,
controlnet,
pipeline,
device,
hf_token,
verbose,
framework_model_dir,
fp16=False,
tf32=False,
bf16=False,
int8=False,
fp8=False,
max_batch_size=16,
do_classifier_free_guidance=False,
):
super(SD3ControlNet, self).__init__(
version,
pipeline,
device=device,
hf_token=hf_token,
verbose=verbose,
framework_model_dir=framework_model_dir,
fp16=fp16,
tf32=tf32,
bf16=bf16,
int8=int8,
fp8=fp8,
max_batch_size=max_batch_size,
)
self.path = load.get_path(version, pipeline, controlnet)
self.subfolder = "controlnet_{}".format(controlnet)
self.controlnet_model_dir = load.get_checkpoint_dir(
self.framework_model_dir, self.version, self.pipeline, self.subfolder
)
self.transformer_model_dir = load.get_checkpoint_dir(
self.framework_model_dir, self.version, self.pipeline, "transformer"
)
if not os.path.exists(self.controlnet_model_dir):
self.config = SD3ControlNetModel.load_config(self.path, token=self.hf_token)
else:
print(f"[I] Load SD3ControlNetModel config from: {self.controlnet_model_dir}")
self.config = SD3ControlNetModel.load_config(self.controlnet_model_dir)
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 {"torch_dtype": torch.bfloat16} if self.bf16 else {}
)
if not load.is_model_cached(self.controlnet_model_dir, model_opts, self.hf_safetensor):
model = SD3ControlNetModel.from_pretrained(self.path, **model_opts, use_safetensors=self.hf_safetensor).to(
self.device
)
model.save_pretrained(self.controlnet_model_dir, **model_opts)
else:
print(f"[I] Load SD3ControlNetModel model from: {self.controlnet_model_dir}")
model = SD3ControlNetModel.from_pretrained(self.controlnet_model_dir, **model_opts).to(self.device)
# Load transformer model for pos_embed
transformer = SD3Transformer2DModel.from_pretrained(self.transformer_model_dir, **model_opts).to(self.device)
if hasattr(model.config, "use_pos_embed") and model.config.use_pos_embed is False:
pos_embed = model._get_pos_embed_from_transformer(transformer)
model.pos_embed = pos_embed.to(model.dtype).to(model.device)
# Free transformer model
del transformer
model = optimizer.optimize_checkpoint(model, torch_inference)
model = SD3ControlNetWrapper(model)
return model
def get_input_names(self):
return ["hidden_states", "controlnet_cond", "conditioning_scale", "pooled_projections", "timestep"]
def get_output_names(self):
return ["controlnet_block_samples"]
def get_dynamic_axes(self):
xB = "2B" if self.xB == 2 else "B"
dynamic_axes = {
"hidden_states": {0: xB, 2: "H", 3: "W"},
"controlnet_cond": {0: xB, 2: "H", 3: "W"},
"pooled_projections": {0: xB},
"timestep": {0: xB},
}
return dynamic_axes
def get_input_profile(
self,
batch_size: int,
image_height: int,
image_width: int,
static_batch: bool,
static_shape: bool,
):
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)
input_profile = {
"hidden_states": [
(self.xB * min_batch, self.config["in_channels"], min_latent_height, min_latent_width),
(self.xB * batch_size, self.config["in_channels"], latent_height, latent_width),
(self.xB * max_batch, self.config["in_channels"], max_latent_height, max_latent_width),
],
"timestep": [(self.xB * min_batch,), (self.xB * batch_size,), (self.xB * max_batch,)],
"pooled_projections": [
(self.xB * min_batch, self.config["pooled_projection_dim"]),
(self.xB * batch_size, self.config["pooled_projection_dim"]),
(self.xB * max_batch, self.config["pooled_projection_dim"]),
],
"controlnet_cond": [
(self.xB * min_batch, self.config["in_channels"], min_latent_height, min_latent_width),
(self.xB * batch_size, self.config["in_channels"], latent_height, latent_width),
(self.xB * max_batch, self.config["in_channels"], max_latent_height, max_latent_width),
],
}
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": (self.xB * batch_size, self.config["in_channels"], latent_height, latent_width),
"timestep": (self.xB * batch_size,),
"pooled_projections": (self.xB * batch_size, self.config["pooled_projection_dim"]),
"controlnet_cond": (self.xB * batch_size, self.config["in_channels"], latent_height, latent_width),
"conditioning_scale": (),
"controlnet_block_samples": (
self.config["num_layers"],
self.xB * batch_size,
latent_height // 2 * latent_width // 2,
self.config["num_attention_heads"] * self.config["attention_head_dim"],
),
}
return shape_dict
def get_sample_input(self, batch_size, image_height, image_width, static_shape):
dtype = torch.float16 if self.fp16 else torch.bfloat16 if self.bf16 else torch.float32
latent_height, latent_width = self.check_dims(batch_size, image_height, image_width)
sample_input = (
torch.randn(
self.xB * batch_size,
self.config["in_channels"],
latent_height,
latent_width,
dtype=dtype,
device=self.device,
),
torch.randn(
self.xB * batch_size,
self.config["in_channels"],
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.config["pooled_projection_dim"], dtype=dtype, device=self.device),
torch.randn(self.xB * batch_size, dtype=torch.float32, device=self.device),
)
return sample_input
demo/Diffusion/demo_diffusion/model/diffusion_transformer.py
+1013
View File
File diff suppressed because it is too large Load Diff
demo/Diffusion/demo_diffusion/model/gan.py
+149
View File
@@ -0,0 +1,149 @@
#
# 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 optimize(self, onnx_graph, return_onnx=True, **kwargs):
onnx_opt_graph = super().optimize(onnx_graph, return_onnx=True, **kwargs)
opt = optimizer.Optimizer(onnx_opt_graph, verbose=self.verbose, version=self.version)
opt.cast_convtranspose_io()
return opt.cleanup(return_onnx=return_onnx)
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
+119
View File
@@ -0,0 +1,119 @@
#
# 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"]
if version == "3.5-large":
return f"stabilityai/stable-diffusion-3.5-large-controlnet-{controlnets}"
return ["lllyasviel/sd-controlnet-" + modality for modality in controlnets]
elif version == "1.4":
return "CompVis/stable-diffusion-v1-4"
elif version == "dreamshaper-7":
return "Lykon/dreamshaper-7"
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 == "3.5-medium":
return "stabilityai/stable-diffusion-3.5-medium"
elif version == "3.5-large":
return "stabilityai/stable-diffusion-3.5-large"
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"
elif version == "flux.1-kontext-dev":
return "black-forest-labs/FLUX.1-Kontext-dev"
elif version == "wan2.2-t2v-a14b":
return "Wan-AI/Wan2.2-T2V-A14B-Diffusers"
elif version == "cosmos-predict2-2b-text2image":
return "nvidia/Cosmos-Predict2-2B-Text2Image"
elif version == "cosmos-predict2-14b-text2image":
return "nvidia/Cosmos-Predict2-14B-Text2Image"
elif version == "cosmos-predict2-2b-video2world":
return "nvidia/Cosmos-Predict2-2B-Video2World"
elif version == "cosmos-predict2-14b-video2world":
return "nvidia/Cosmos-Predict2-14B-Video2World"
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
+56
View File
@@ -0,0 +1,56 @@
#
# 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 abc import ABC
from diffusers.loaders import StableDiffusionLoraLoaderMixin, FluxLoraLoaderMixin
class LoraLoader(ABC):
def __init__(self, paths, weights, scale):
self.paths = paths
self.weights = weights
self.scale = scale
class SDLoraLoader(LoraLoader, StableDiffusionLoraLoaderMixin):
def __init__(self, paths, weights, scale):
super().__init__(paths, weights, scale)
class FLUXLoraLoader(LoraLoader, FluxLoraLoaderMixin):
def __init__(self, paths, weights, scale):
super().__init__(paths, weights, 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]}")
if isinstance(lora_loader, SDLoraLoader):
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)
elif isinstance(lora_loader, FLUXLoraLoader):
state_dict, network_alphas = lora_loader.lora_state_dict(path, return_alphas=True)
lora_loader.load_lora_into_transformer(state_dict, network_alphas=network_alphas, transformer=model, adapter_name=path)
else:
raise ValueError(f"Unsupported LoRA loader: {lora_loader}")
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
+218
View File
@@ -0,0 +1,218 @@
#
# 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
import onnx_graphsurgeon as gs
import torch
from onnx import shape_inference
from polygraphy.backend.onnx.loader import fold_constants
from demo_diffusion.model import load
from demo_diffusion.utils_modelopt import (
cast_convtranspose_io,
cast_fp8_mha_io,
cast_layernorm_io,
cast_resize_io,
convert_fp16_io,
convert_zp_fp8,
)
# 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, version=None):
self.graph = gs.import_onnx(onnx_graph)
self.verbose = verbose
self.version = version
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_int8_graph(self):
# Cast LayerNorm scale/bias from FP16 to FP32 to match INT8 DQ activations.
cast_layernorm_io(self.graph)
# Fuse QKV QDQ nodes for INT8 SmoothQuant.
self.fuse_mha_qkv_int8_sq()
def cast_convtranspose_io(self):
cast_convtranspose_io(self.graph)
def cast_resize_io(self, output_dtype):
cast_resize_io(self.graph, output_dtype=output_dtype)
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)
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
+142
View File
@@ -0,0 +1,142 @@
#
# 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,
UMT5EncoderModel,
)
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,
weight_streaming=False,
weight_streaming_budget_percentage=None,
use_attention_mask=False,
):
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.is_umt5 = getattr(self.config, 'model_type', '') == 'umt5'
self.weight_streaming = weight_streaming
self.weight_streaming_budget_percentage = weight_streaming_budget_percentage
self.use_attention_mask = use_attention_mask
def get_model(self, torch_inference=""):
model_opts = (
{"torch_dtype": torch.float16} if self.fp16 else {"torch_dtype": torch.bfloat16} if self.bf16 else {}
)
EncoderModelClass = UMT5EncoderModel if self.is_umt5 else T5EncoderModel
if not load.is_model_cached(self.t5_model_dir, model_opts, self.hf_safetensor, model_name="model"):
model = EncoderModelClass.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 {EncoderModelClass.__name__} model from: {self.t5_model_dir}")
model = EncoderModelClass.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):
if self.use_attention_mask:
return ["input_ids", "attention_mask"]
return ["input_ids"]
def get_output_names(self):
return ["text_embeddings"]
def get_dynamic_axes(self):
if self.use_attention_mask:
return {"input_ids": {0: "B"}, "attention_mask": {0: "B"}, "text_embeddings": {0: "B"}}
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
)
profile = {
"input_ids": [(min_batch, self.text_maxlen), (batch_size, self.text_maxlen), (max_batch, self.text_maxlen)]
}
if self.use_attention_mask:
profile["attention_mask"] = [
(min_batch, self.text_maxlen),
(batch_size, self.text_maxlen),
(max_batch, self.text_maxlen),
]
return profile
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),
}
if self.use_attention_mask:
output["attention_mask"] = (batch_size, self.text_maxlen)
return output
def get_sample_input(self, batch_size, image_height, image_width, static_shape):
self.check_dims(batch_size, image_height, image_width)
inputs = {"input_ids": torch.zeros(batch_size, self.text_maxlen, dtype=torch.int32, device=self.device)}
if self.use_attention_mask:
inputs["attention_mask"] = torch.ones(batch_size, self.text_maxlen, dtype=torch.int32, device=self.device)
return inputs
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
+902
View File
@@ -0,0 +1,902 @@
#
# 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
from diffusers import StableDiffusionXLControlNetPipeline
# 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", "dreamshaper-7"):
return 768
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 = 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, modify_int8_graph=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 = 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, modify_int8_graph=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_pipeline(self):
cnet_model_opts = {"torch_dtype": torch.float16} if self.fp16 else {}
controlnets = [
ControlNetModel.from_pretrained(path, **cnet_model_opts).to(self.device) for path in self.controlnets
]
if self.bf16:
model_opts = {"torch_dtype": torch.bfloat16}
elif self.fp16:
model_opts = {"variant": "fp16", "torch_dtype": torch.float16}
else:
model_opts = {}
pipeline = StableDiffusionXLControlNetPipeline.from_pretrained(
self.path,
use_safetensors=self.hf_safetensor,
token=self.hf_token,
controlnet=controlnets,
**model_opts,
).to(self.device)
return pipeline
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
+671
View File
@@ -0,0 +1,671 @@
#
# 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", "AutoencoderKLWan"]
for model in models_to_import:
globals()[model] = import_from_diffusers(model, "diffusers.models")
# Import FluxKontextUtil from pipeline module
# Using a deferred import to avoid circular dependencies
def _get_flux_kontext_util():
from demo_diffusion.pipeline.flux_pipeline import FluxKontextUtil
return FluxKontextUtil
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.do_constant_folding = False
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,
do_classifier_free_guidance=False,
kontext_resolution=None,
):
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.kontext_resolution = kontext_resolution
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)
self.xB = 2 if do_classifier_free_guidance else 1 # batch multiplier
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):
xB = "2B" if self.xB == 2 else "B"
return {"images": {0: xB, 2: "8H", 3: "8W"}, "latent": {0: xB, 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)
)
if self.version == "flux.1-kontext-dev":
FluxKontextUtil = _get_flux_kontext_util()
min_latent_dim, max_latent_dim = FluxKontextUtil.get_min_max_kontext_dimensions()
return {
"images": [
(self.xB * min_batch, 3, min_latent_dim[1], min_latent_dim[0]),
(self.xB * batch_size, 3, self.kontext_resolution[1], self.kontext_resolution[0]),
(self.xB * max_batch, 3, max_latent_dim[1], max_latent_dim[0]),
],
}
return {
"images": [
(self.xB * min_batch, 3, min_image_height, min_image_width),
(self.xB * batch_size, 3, image_height, image_width),
(self.xB * max_batch, 3, max_image_height, max_image_width),
],
}
def get_shape_dict(self, batch_size, image_height, image_width):
# Determine dimensions based on version
if self.version == "flux.1-kontext-dev":
img_h, img_w = self.kontext_resolution[1], self.kontext_resolution[0]
else:
img_h, img_w = image_height, image_width
latent_height, latent_width = self.check_dims(batch_size, img_h, img_w)
return {
"images": (self.xB * batch_size, 3, img_h, img_w),
"latent": (self.xB * 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(self.xB * batch_size, 3, image_height, image_width, dtype=dtype, device=self.device)
class SD3_VAEEncoderModel(base_model.BaseModel):
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_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):
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)
class AutoencoderKLWanModel(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(AutoencoderKLWanModel, 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.is_wan_pipeline = version.startswith("wan")
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 = AutoencoderKLWan.load_config(self.path, subfolder=self.subfolder, token=self.hf_token)
else:
print(f"[I] Load AutoencoderKLWan (decoder) config from: {self.vae_decoder_model_dir}")
self.config = AutoencoderKLWan.load_config(self.vae_decoder_model_dir)
def get_model(self, torch_inference=""):
if self.is_wan_pipeline:
print(f"[I] Using float32 precision for Wan 2.2 VAE decoder")
model_opts = {"torch_dtype": torch.float32}
else:
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 = AutoencoderKLWan.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 AutoencoderKLWan (decoder) model from: {self.vae_decoder_model_dir}")
model = AutoencoderKLWan.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", 3: "H", 4: "W"}, "images": {0: "B", 3: "8H", 4: "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["z_dim"], 1, min_latent_height, min_latent_width),
(batch_size, self.config["z_dim"], 1, latent_height, latent_width),
(max_batch, self.config["z_dim"], 1, 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["z_dim"], 1, latent_height, latent_width),
"images": (batch_size, 3, 1, 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["z_dim"], 1, latent_height, latent_width, dtype=dtype, device=self.device
)
class AutoencoderKLWanEncoderModelWrapper(torch.nn.Module):
def __init__(self, model):
super().__init__()
self.model = model
def forward(self, x):
return self.model.encode(x).latent_dist.sample()
class AutoencoderKLWanEncoderModel(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(AutoencoderKLWanEncoderModel, 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(
self.framework_model_dir, self.version, self.pipeline, self.subfolder
)
if not os.path.exists(self.vae_encoder_model_dir):
self.config = AutoencoderKLWan.load_config(self.path, subfolder=self.subfolder, token=self.hf_token)
else:
print(f"[I] Load AutoencoderKLWan (encoder) config from: {self.vae_encoder_model_dir}")
self.config = AutoencoderKLWan.load_config(self.vae_encoder_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_encoder_model_dir, model_opts, self.hf_safetensor):
model = AutoencoderKLWan.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_encoder_model_dir, **model_opts)
else:
print(f"[I] Load AutoencoderKLWan (encoder) model from: {self.vae_encoder_model_dir}")
model = AutoencoderKLWan.from_pretrained(self.vae_encoder_model_dir, **model_opts).to(self.device)
model = AutoencoderKLWanEncoderModelWrapper(model)
model = optimizer.optimize_checkpoint(model, torch_inference)
return model
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
+50
View File
@@ -0,0 +1,50 @@
#
# 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."""
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
+168
View File
@@ -0,0 +1,168 @@
#
# 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"
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():
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
+104
View File
@@ -0,0 +1,104 @@
#
# 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 importlib
from typing import TYPE_CHECKING
# Expose public API while avoiding importing optional dependencies at module import time.
# Each attribute is imported on first access via __getattr__.
__all__ = [
"DiffusionPipeline",
"FluxPipeline",
"FluxKontextPipeline",
"StableCascadePipeline",
"StableDiffusion3Pipeline",
"StableDiffusion35Pipeline",
"StableDiffusionPipeline",
"CosmosPipeline",
"StableVideoDiffusionPipeline",
"WanPipeline",
"PIPELINE_TYPE",
]
_LAZY_ATTRS = {
# Core/base
"DiffusionPipeline": ("demo_diffusion.pipeline.diffusion_pipeline", "DiffusionPipeline"),
"PIPELINE_TYPE": ("demo_diffusion.pipeline.type", "PIPELINE_TYPE"),
# Stable Diffusion family
"StableDiffusionPipeline": ("demo_diffusion.pipeline.stable_diffusion_pipeline", "StableDiffusionPipeline"),
"StableDiffusion3Pipeline": ("demo_diffusion.pipeline.stable_diffusion_3_pipeline", "StableDiffusion3Pipeline"),
"StableDiffusion35Pipeline": ("demo_diffusion.pipeline.stable_diffusion_35_pipeline", "StableDiffusion35Pipeline"),
# Stable Cascade
"StableCascadePipeline": ("demo_diffusion.pipeline.stable_cascade_pipeline", "StableCascadePipeline"),
# Stable Video Diffusion
"StableVideoDiffusionPipeline": ("demo_diffusion.pipeline.stable_video_diffusion_pipeline", "StableVideoDiffusionPipeline"),
# Flux family (optional dependency: `flux`)
"FluxPipeline": ("demo_diffusion.pipeline.flux_pipeline", "FluxPipeline"),
"FluxKontextPipeline": ("demo_diffusion.pipeline.flux_pipeline", "FluxKontextPipeline"),
# Cosmos (optional dependency: `flux`)
"CosmosPipeline": ("demo_diffusion.pipeline.cosmos_pipeline", "CosmosPipeline"),
# Wan
"WanPipeline": ("demo_diffusion.pipeline.wan_pipeline", "WanPipeline"),
}
def __getattr__(name):
if name not in _LAZY_ATTRS:
raise AttributeError(f"module 'demo_diffusion.pipeline' has no attribute {name!r}")
module_path, attr_name = _LAZY_ATTRS[name]
try:
module = importlib.import_module(module_path)
except ModuleNotFoundError as e:
missing_pkg = e.name or "<unknown>"
raise ModuleNotFoundError(
f"Optional dependency '{missing_pkg}' is required for '{name}'. "
"Install the appropriate extras/requirements for the selected pipeline "
"(e.g., use the non-legacy requirements for Flux/Cosmos), or install the missing package."
) from e
try:
return getattr(module, attr_name)
except AttributeError as e:
raise AttributeError(
f"'{module_path}' does not export attribute '{attr_name}' (while resolving '{name}')."
) from e
if TYPE_CHECKING:
from demo_diffusion.pipeline.diffusion_pipeline import DiffusionPipeline as DiffusionPipeline
from demo_diffusion.pipeline.type import PIPELINE_TYPE as PIPELINE_TYPE
from demo_diffusion.pipeline.stable_diffusion_pipeline import (
StableDiffusionPipeline as StableDiffusionPipeline,
)
from demo_diffusion.pipeline.stable_diffusion_3_pipeline import (
StableDiffusion3Pipeline as StableDiffusion3Pipeline,
)
from demo_diffusion.pipeline.stable_diffusion_35_pipeline import (
StableDiffusion35Pipeline as StableDiffusion35Pipeline,
)
from demo_diffusion.pipeline.stable_cascade_pipeline import (
StableCascadePipeline as StableCascadePipeline,
)
from demo_diffusion.pipeline.stable_video_diffusion_pipeline import (
StableVideoDiffusionPipeline as StableVideoDiffusionPipeline,
)
# Optional pipelines
from demo_diffusion.pipeline.flux_pipeline import (
FluxPipeline as FluxPipeline,
FluxKontextPipeline as FluxKontextPipeline,
)
from demo_diffusion.pipeline.cosmos_pipeline import CosmosPipeline as CosmosPipeline
from demo_diffusion.pipeline.wan_pipeline import WanPipeline as WanPipeline
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/demo_diffusion/pipeline/cosmos_pipeline.py
+1007
View File
File diff suppressed because it is too large Load Diff
demo/Diffusion/demo_diffusion/pipeline/diffusion_pipeline.py
Executable
+997
View File
@@ -0,0 +1,997 @@
#
# 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 abc
import argparse
import gc
import json
import os
import pathlib
import sys
from abc import ABC, abstractmethod
from typing import Any, List
import modelopt.torch.opt as mto
import modelopt.torch.quantization as mtq
import nvtx
import torch
from cuda.bindings import runtime as cudart
from diffusers import (
DDIMScheduler,
DDPMScheduler,
DDPMWuerstchenScheduler,
EulerAncestralDiscreteScheduler,
EulerDiscreteScheduler,
FlowMatchEulerDiscreteScheduler,
LCMScheduler,
LMSDiscreteScheduler,
PNDMScheduler,
UniPCMultistepScheduler,
)
from torch.utils.data import DataLoader
import demo_diffusion.engine as engine_module
import demo_diffusion.image as image_module
from demo_diffusion.model import (
make_scheduler,
merge_loras,
unload_torch_model,
)
from demo_diffusion.pipeline.calibrate import load_calib_prompts
from demo_diffusion.pipeline.model_memory_manager import ModelMemoryManager
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,
PromptImageDataset,
SameSizeSampler,
check_lora,
custom_collate,
filter_func,
filter_func_no_proj_out,
fp8_mha_disable,
generate_fp8_scales,
get_int8_config,
infinite_dataloader,
quantize_lvl,
set_fmha,
set_quant_precision,
)
class DiffusionPipeline(ABC):
"""
Application showcasing the acceleration of Stable Diffusion pipelines using NVidia TensorRT.
"""
VALID_DIFFUSION_PIPELINES = (
"1.4",
"dreamshaper-7",
"xl-1.0",
"xl-turbo",
"svd-xt-1.1",
"sd3",
"3.5-medium",
"3.5-large",
"cascade",
"flux.1-dev",
"flux.1-dev-canny",
"flux.1-dev-depth",
"flux.1-schnell",
"flux.1-kontext-dev",
"wan2.2-t2v-a14b",
"cosmos-predict2-2b-text2image",
"cosmos-predict2-14b-text2image",
"cosmos-predict2-2b-video2world",
"cosmos-predict2-14b-video2world",
)
SCHEDULER_DEFAULTS = {
"1.4": "PNDM",
"dreamshaper-7": "PNDM",
"xl-1.0": "Euler",
"xl-turbo": "EulerA",
"3.5-large": "FlowMatchEuler",
"3.5-medium": "FlowMatchEuler",
"svd-xt-1.1": "Euler",
"cascade": "DDPMWuerstchen",
"flux.1-dev": "FlowMatchEuler",
"flux.1-dev-canny": "FlowMatchEuler",
"flux.1-dev-depth": "FlowMatchEuler",
"flux.1-schnell": "FlowMatchEuler",
"flux.1-kontext-dev": "FlowMatchEuler",
"wan2.2-t2v-a14b": "UniPC",
"cosmos-predict2-2b-text2image": "FlowMatchEuler",
"cosmos-predict2-14b-text2image": "FlowMatchEuler",
"cosmos-predict2-2b-video2world": "FlowMatchEuler",
"cosmos-predict2-14b-video2world": "FlowMatchEuler",
}
def __init__(
self,
dd_path,
version="1.4",
pipeline_type=PIPELINE_TYPE.TXT2IMG,
bf16=False,
max_batch_size=16,
denoising_steps=30,
scheduler=None,
device="cuda",
output_dir=".",
hf_token=None,
verbose=False,
nvtx_profile=False,
use_cuda_graph=False,
framework_model_dir="pytorch_model",
return_latents=False,
low_vram=False,
torch_inference="",
torch_fallback=None,
weight_streaming=False,
text_encoder_weight_streaming_budget_percentage=None,
denoiser_weight_streaming_budget_percentage=None,
controlnet=None,
):
"""
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):
Task performed by the current pipeline. Should be one of PIPELINE_TYPE.__members__.
max_batch_size (int):
Maximum batch size for dynamic batch engine.
bf16 (`bool`, defaults to False):
Whether to run the pipeline in BFloat16 precision.
denoising_steps (int):
The number of denoising steps.
More denoising steps usually lead to a higher quality image at the expense of slower inference.
scheduler (str):
The scheduler to guide the denoising process. Must be one of the values listed in DiffusionPipeline.SCHEDULER_DEFAULTS.values().
device (str):
PyTorch device to run inference. Default: 'cuda'.
output_dir (str):
Output directory for log files and image artifacts.
hf_token (str):
HuggingFace User Access Token to use for downloading Stable Diffusion model checkpoints.
verbose (bool):
Enable verbose logging.
nvtx_profile (bool):
Insert NVTX profiling markers.
use_cuda_graph (bool):
Use CUDA graph to capture engine execution and then launch inference.
framework_model_dir (str):
cache directory for framework checkpoints.
return_latents (bool):
Skip decoding the image and return latents instead.
low_vram (bool):
[FLUX only] Optimize for low VRAM usage, possibly at the expense of inference performance. Disabled by default.
torch_inference (str):
Run inference with PyTorch (using specified compilation mode) instead of TensorRT. The compilation mode specified should be one of ['eager', 'reduce-overhead', 'max-autotune'].
torch_fallback (str):
[FLUX only] Comma separated list of models to be inferenced using PyTorch instead of TRT. For example --torch-fallback t5,transformer. If --torch-inference set, this parameter will be ignored.
weight_streaming (`bool`, defaults to False):
Whether to enable weight streaming during TensorRT engine build.
text_encoder_ws_budget_percentage (`int`, defaults to None):
Weight streaming budget as a percentage of the size of total streamable weights for the text encoder model.
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.
controlnet (str, defaults to None):
Type of ControlNet to use for the pipeline.
"""
self.bf16 = bf16
self.dd_path = dd_path
self.denoising_steps = denoising_steps
self.max_batch_size = max_batch_size
self.framework_model_dir = framework_model_dir
self.output_dir = output_dir
for directory in [self.framework_model_dir, self.output_dir]:
if not os.path.exists(directory):
print(f"[I] Create directory: {directory}")
pathlib.Path(directory).mkdir(parents=True)
self.hf_token = hf_token
self.device = device
self.verbose = verbose
self.nvtx_profile = nvtx_profile
self.version = version
self.pipeline_type = pipeline_type
self.return_latents = return_latents
self.low_vram = low_vram
self.weight_streaming = weight_streaming
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, controlnet)
# config to store additional info
self.config = {}
if torch_fallback:
assert type(torch_fallback) is list
for model_name in torch_fallback:
if model_name not in self.stages:
raise ValueError(f'Model "{model_name}" set in --torch-fallback does not exist')
self.config[model_name.replace("-", "_") + "_torch_fallback"] = True
print(f"[I] Setting torch_fallback for {model_name} model.")
if not scheduler:
scheduler = 'UniPC' if self.pipeline_type.is_controlnet() and not self.version == "3.5-large" else self.SCHEDULER_DEFAULTS.get(version, 'DDIM')
print(f"[I] Autoselected scheduler: {scheduler}")
scheduler_class_map = {
"DDIM" : DDIMScheduler,
"DDPM" : DDPMScheduler,
"EulerA" : EulerAncestralDiscreteScheduler,
"Euler" : EulerDiscreteScheduler,
"LCM" : LCMScheduler,
"LMSD" : LMSDiscreteScheduler,
"PNDM" : PNDMScheduler,
"UniPC" : UniPCMultistepScheduler,
"DDPMWuerstchen" : DDPMWuerstchenScheduler,
"FlowMatchEuler": FlowMatchEulerDiscreteScheduler,
}
try:
scheduler_class = scheduler_class_map[scheduler]
except KeyError:
raise ValueError(
f"Unsupported scheduler {scheduler}. Should be one of {list(scheduler_class_map.keys())}."
)
self.scheduler = make_scheduler(scheduler_class, version, pipeline_type, hf_token, framework_model_dir)
self.torch_inference = torch_inference
if self.torch_inference:
torch._inductor.config.conv_1x1_as_mm = True
torch._inductor.config.coordinate_descent_tuning = True
torch._inductor.config.epilogue_fusion = False
torch._inductor.config.coordinate_descent_check_all_directions = True
self.use_cuda_graph = use_cuda_graph
# initialized in load_engines()
self.models = {}
self.torch_models = {}
self.engine = {}
self.shape_dicts = {}
self.shared_device_memory = None
self.lora_loader = None
# initialized in load_resources()
self.events = {}
self.generator = None
self.markers = {}
self.seed = None
self.stream = None
self.tokenizer = None
def model_memory_manager(self, model_names, low_vram=False):
return ModelMemoryManager(self, model_names, low_vram)
@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, controlnet_type: str = None) -> 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, 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__}_{version}"
def profile_start(self, name, color="blue", domain=None):
if self.nvtx_profile:
self.markers[name] = nvtx.start_range(message=name, color=color, domain=domain)
if name in self.events:
cudart.cudaEventRecord(self.events[name][0], 0)
def profile_stop(self, name):
if name in self.events:
cudart.cudaEventRecord(self.events[name][1], 0)
if self.nvtx_profile:
nvtx.end_range(self.markers[name])
def load_resources(self, image_height, image_width, batch_size, seed):
# Initialize noise generator
if seed is not None:
self.seed = seed
self.generator = torch.Generator(device="cuda").manual_seed(seed)
# Create CUDA events and stream
for stage in self.stages:
self.events[stage] = [cudart.cudaEventCreate()[1], cudart.cudaEventCreate()[1]]
self.stream = cudart.cudaStreamCreate()[1]
# Allocate TensorRT I/O buffers
if not self.torch_inference:
for model_name, obj in self.models.items():
if self.torch_fallback[model_name]:
continue
shape_dict = obj.get_shape_dict(
batch_size, image_height, image_width,
**({'num_frames': obj.num_frames} if hasattr(obj, 'num_frames') and obj.num_frames else {})
)
self.shape_dicts[model_name] = shape_dict
if not self.low_vram:
self.engine[model_name].allocate_buffers(shape_dict=shape_dict, device=self.device)
@abstractmethod
def _initialize_models(self, *args, **kwargs):
raise NotImplementedError("Please Implement the _initialize_models method")
def _prepare_model_configs(
self,
enable_refit,
int8,
fp8,
fp4
):
model_names = self.models.keys()
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:
# Initialize config
do_engine_refit = enable_refit and not self.pipeline_type.is_sd_xl_refiner() and any(model_name.startswith(prefix) for prefix in ("unet", "transformer"))
do_lora_merge = not enable_refit and self.lora_loader and any(model_name.startswith(prefix) for prefix in ("unet", "transformer"))
config = {
"do_engine_refit": do_engine_refit,
"do_lora_merge": do_lora_merge,
"use_int8": False,
"use_fp8": False,
'use_fp4': False,
}
# TODO: Move this to when arguments are first being validated in dd_argparse.py
# 8-bit/4-bit precision inference
if int8:
assert self.pipeline_type.is_sd_xl_base() or self.version in [
"1.4",
], "int8 quantization only supported for SDXL and SD1.4 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.4"]
or self.version.startswith("flux.1")
or self.version.startswith("3.5-large")
), "fp8 quantization only supported for SDXL, SD1.4, SD3.5-large and FLUX pipelines"
if (
(self.pipeline_type.is_sd_xl() and model_name == "unetxl")
or (self.version.startswith("flux.1") and model_name == "transformer")
or (
self.version.startswith("3.5-large")
and ("transformer" in model_name or "controlnet" in model_name)
)
or (model_name == "unet")
):
config["use_fp8"] = True
elif fp4:
config['use_fp4'] = True
# Setup paths
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
return configs
def _calibrate_and_save_model(
self,
pipeline,
model,
model_config,
quantization_level,
quantization_percentile,
quantization_alpha,
calibration_size,
calib_batch_size,
enable_lora_merge = False,
**kwargs):
print(f"[I] Calibrated weights not found, generating {model_config['state_dict_path']}")
# TODO check size > calibration_size
def do_calibrate(pipeline, calibration_prompts, **kwargs):
for i_th, prompts in enumerate(calibration_prompts):
if i_th >= kwargs["calib_size"]:
return
if kwargs["model_id"] in ("flux.1-dev", "flux.1-schnell"):
common_args = {
"prompt": prompts,
"prompt_2": prompts,
"num_inference_steps": kwargs["n_steps"],
"height": kwargs.get("height", 1024),
"width": kwargs.get("width", 1024),
"guidance_scale": 3.5,
"max_sequence_length": 512 if kwargs["model_id"] == "flux.1-dev" else 256,
}
else:
common_args = {
"prompt": prompts,
"num_inference_steps": kwargs["n_steps"],
"negative_prompt": ["normal quality, low quality, worst quality, low res, blurry, nsfw, nude"]
* len(prompts),
}
pipeline(**common_args).images
def do_calibrate_img2img(pipeline, dataloader, **kwargs):
for i_th, (img_conds, prompts) in enumerate(dataloader):
if i_th >= kwargs["calib_size"]:
return
common_args = {
"prompt": list(prompts),
"control_image": img_conds,
"num_inference_steps": kwargs["n_steps"],
"height": img_conds.size(2),
"width": img_conds.size(3),
"generator": torch.Generator().manual_seed(42),
"guidance_scale": 3.5,
"max_sequence_length": 512,
}
pipeline(**common_args).images
if self.version in ("flux.1-dev-depth", "flux.1-dev-canny"):
dataset = PromptImageDataset(
root_dir=self.calibration_dataset,
)
dataloader = DataLoader(
dataset,
batch_size=calib_batch_size,
shuffle=False,
num_workers=0,
sampler=SameSizeSampler(dataset=dataset, batch_size=calib_batch_size),
collate_fn=custom_collate,
)
else:
root_dir = os.path.dirname(os.path.abspath(sys.modules["__main__"].__file__))
calibration_file = os.path.join(root_dir, "calibration_data", "calibration-prompts.txt")
calibration_prompts = load_calib_prompts(calib_batch_size, calibration_file)
def forward_loop(model):
if self.version not in ("sd3", "flux.1-dev", "flux.1-schnell", "flux.1-dev-depth", "flux.1-dev-canny"):
pipeline.unet = model
else:
pipeline.transformer = model
if self.version in ("flux.1-dev-depth", "flux.1-dev-canny"):
do_calibrate_img2img(
pipeline=pipeline,
dataloader=infinite_dataloader(dataloader),
calib_size=calibration_size // calib_batch_size,
n_steps=self.denoising_steps,
model_id=self.version,
)
else:
do_calibrate(
pipeline=pipeline,
calibration_prompts=calibration_prompts,
calib_size=calibration_size // calib_batch_size,
n_steps=self.denoising_steps,
model_id=self.version,
**kwargs
)
print(f"[I] Performing calibration for {calibration_size} steps.")
if model_config['use_int8']:
quant_config = get_int8_config(
model,
quantization_level,
quantization_alpha,
quantization_percentile,
self.denoising_steps
)
elif model_config['use_fp8']:
if self.version.startswith("flux.1"):
quant_config = SD_FP8_BF16_FLUX_MMDIT_BMM2_FP8_OUTPUT_CONFIG
else:
quant_config = SD_FP8_FP16_DEFAULT_CONFIG
# Handle LoRA
if enable_lora_merge:
assert self.lora_loader is not None
model = merge_loras(model, self.lora_loader)
check_lora(model)
if self.version.startswith("flux.1"):
set_quant_precision(quant_config, "BFloat16")
mtq.quantize(model, quant_config, forward_loop)
mto.save(model, model_config['state_dict_path'])
def _get_quantized_model(
self,
obj,
model_config,
quantization_level,
quantization_percentile,
quantization_alpha,
calibration_size,
calib_batch_size,
enable_lora_merge = False,
**kwargs):
pipeline = obj.get_pipeline()
is_flux = self.version.startswith("flux.1")
model = pipeline.unet if self.version not in ("sd3", "flux.1-dev", "flux.1-schnell", "flux.1-dev-depth", "flux.1-dev-canny") else pipeline.transformer
if model_config['use_fp8'] and quantization_level == 4.0:
set_fmha(model, is_flux=is_flux)
if not os.path.exists(model_config['state_dict_path']):
self._calibrate_and_save_model(
pipeline,
model,
model_config,
quantization_level,
quantization_percentile,
quantization_alpha,
calibration_size,
calib_batch_size,
enable_lora_merge,
**kwargs)
else:
mto.restore(model, model_config['state_dict_path'])
if not os.path.exists(model_config['onnx_path']):
quantize_lvl(self.version, model, quantization_level)
if self.version.startswith("flux.1"):
mtq.disable_quantizer(model, filter_func_no_proj_out)
else:
mtq.disable_quantizer(model, filter_func)
if model_config['use_fp8'] and not self.version.startswith("flux.1"):
generate_fp8_scales(model)
if quantization_level == 4.0:
fp8_mha_disable(model, quantized_mha_output=False) # Remove Q/DQ after BMM2 in MHA
else:
model = None
return model
@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
@staticmethod
def _fix_bf16_resize_nodes(onnx_opt_path):
"""Cast Resize node I/O for strongly-typed TRT engine builds.
TRT does not support BF16 for the Resize operator, so inputs are
cast to FP32 and outputs are cast back to BF16."""
import onnx
import onnx_graphsurgeon as gs
from demo_diffusion.model import load
from demo_diffusion.utils_modelopt import cast_resize_io
onnx_graph = onnx.load(onnx_opt_path, load_external_data=True)
graph = gs.import_onnx(onnx_graph)
resize_nodes = [n for n in graph.nodes if n.op == "Resize"]
if not resize_nodes:
return
print(f"[I] Fixing {len(resize_nodes)} BF16 Resize node(s) in downloaded model: {onnx_opt_path}")
cast_resize_io(graph, output_dtype=onnx.TensorProto.BFLOAT16)
graph.cleanup().toposort()
onnx_graph = gs.export_onnx(graph)
if load.onnx_graph_needs_external_data(onnx_graph):
onnx.save_model(
onnx_graph,
onnx_opt_path,
save_as_external_data=True,
all_tensors_to_one_file=True,
convert_attribute=False,
)
else:
onnx.save(onnx_graph, onnx_opt_path)
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 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)
# Fix Resize nodes for strongly-typed TRT builds.
# Downloaded models bypass optimize(), so apply the fix here.
if obj.bf16:
self._fix_bf16_resize_nodes(model_config['onnx_opt_path'])
do_export_onnx = False
else:
self.is_native_export_supported(model_config)
dynamo = True if (self.pipeline_type.is_video2world() and model_name == "transformer") or (self.pipeline_type.is_txt2vid() and (model_name in ["transformer", "transformer_2"])) or (self.version.startswith("flux.1") and model_name == "transformer" and obj.fp16) else False
export_kwargs = {
"static_shape": static_shape,
"dynamo": dynamo,
**({'opt_num_frames': obj.num_frames} if hasattr(obj, 'num_frames') and obj.num_frames else {})
}
if do_export_onnx or do_export_weights_map:
if not model_config['use_int8'] and not model_config['use_fp8']:
obj.export_onnx(
model_config["onnx_path"],
model_config["onnx_opt_path"],
onnx_opset,
opt_image_height,
opt_image_width,
enable_lora_merge=model_config["do_lora_merge"],
lora_loader=self.lora_loader,
**export_kwargs,
)
else:
print(f"[I] Generating quantized ONNX model: {model_config['onnx_path']}")
quantized_model = self._get_quantized_model(
obj,
model_config,
quantization_level,
quantization_percentile,
quantization_alpha,
calibration_size,
calib_batch_size,
height=opt_image_width,
width=opt_image_width,
enable_lora_merge=model_config["do_lora_merge"],
)
obj.export_onnx(
model_config["onnx_path"],
model_config["onnx_opt_path"],
onnx_opset,
opt_image_height,
opt_image_width,
custom_model=quantized_model,
**export_kwargs,
)
# FIXME do_export_weights_map needs ONNX graph
if do_export_weights_map:
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
tf32amp = obj.tf32
weight_streaming = getattr(obj, 'weight_streaming', False)
precision_constraints = 'none'
input_profile = obj.get_input_profile(
opt_batch_size, opt_image_height, opt_image_width,
static_batch=static_batch, static_shape=static_shape,
**({'num_frames': obj.num_frames} if hasattr(obj, 'num_frames') and obj.num_frames else {})
)
engine.build(
model_config["onnx_opt_path"],
tf32=tf32amp,
input_profile=input_profile,
enable_refit=model_config["do_engine_refit"],
enable_all_tactics=enable_all_tactics,
timing_cache=timing_cache,
update_output_names=update_output_names,
weight_streaming=weight_streaming,
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']
with open(model_config['weights_map_path'], 'r') as fp_wts:
print(f"[I] Loading weights map: {model_config['weights_map_path']} ")
[weights_name_mapping, weights_shape_mapping] = json.load(fp_wts)
if not os.path.exists(model_config['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(), 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_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'])
self.engine[model_name].refit(refit_weights, updated_weight_names)
def _load_torch_models(self):
# Load torch models
for model_name, obj in self.models.items():
if self.torch_fallback[model_name]:
self.torch_models[model_name] = obj.get_model(torch_inference=self.torch_inference)
if self.low_vram:
self.torch_models[model_name] = self.torch_models[model_name].to('cpu')
torch.cuda.empty_cache()
def load_engines(
self,
framework_model_dir,
onnx_opset,
opt_batch_size,
opt_image_height,
opt_image_width,
optimization_level=3,
static_batch=False,
static_shape=True,
enable_refit=False,
enable_all_tactics=False,
timing_cache=None,
int8=False,
fp8=False,
fp4=False,
quantization_level=2.5,
quantization_percentile=1.0,
quantization_alpha=0.8,
calibration_size=32,
calib_batch_size=2,
onnx_export_only=False,
download_onnx_models=False,
):
"""
Build and load engines for TensorRT accelerated inference.
Export ONNX models first, if applicable.
Args:
framework_model_dir (str):
Directory to store the framework model ckpt.
onnx_opset (int):
ONNX opset version to export the models.
opt_batch_size (int):
Batch size to optimize for during engine building.
opt_image_height (int):
Image height to optimize for during engine building. Must be a multiple of 8.
opt_image_width (int):
Image width to optimize for during engine building. Must be a multiple of 8.
optimization_level (int):
Optimization level to build the TensorRT engine with.
static_batch (bool):
Build engine only for specified opt_batch_size.
static_shape (bool):
Build engine only for specified opt_image_height & opt_image_width. Default = True.
enable_refit (bool):
Build engines with refit option enabled.
enable_all_tactics (bool):
Enable all tactic sources during TensorRT engine builds.
timing_cache (str):
Path to the timing cache to speed up TensorRT build.
int8 (bool):
Whether to quantize to int8 format or not (SDXL, SD15 and SD21 only).
fp8 (bool):
Whether to quantize to fp8 format or not (SDXL, SD15 and SD21 only).
quantization_level (float):
Controls which layers to quantize. 1: CNN, 2: CNN+FFN, 2.5: CNN+FFN+QKV, 3: CNN+FC
quantization_percentile (float):
Control quantization scaling factors (amax) collecting range, where the minimum amax in
range(n_steps * percentile) will be collected. Recommendation: 1.0
quantization_alpha (float):
The alpha parameter for SmoothQuant quantization used for linear layers.
Recommendation: 0.8 for SDXL
calibration_size (int):
The number of steps to use for calibrating the model for quantization.
Recommendation: 32, 64, 128 for SDXL
calib_batch_size (int):
The batch size to use for calibration. Defaults to 2.
onnx_export_only (bool):
Whether only export onnx without building the TRT engine.
download_onnx_models (bool):
Download pre-exported ONNX models
"""
self._initialize_models(framework_model_dir, int8, fp8, fp4)
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_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()
torch.cuda.empty_cache()
if onnx_export_only:
return
# Build TensorRT engines
for model_name, obj in self.models.items():
if self.torch_fallback[model_name]:
continue
model_config = model_configs[model_name]
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
# Load and refit TensorRT engines
for model_name, obj in self.models.items():
if self.torch_fallback[model_name]:
continue
model_config = model_configs[model_name]
# For non low_vram case, the engines will remain in GPU memory from now on.
assert self.engine[model_name].engine is None
if not self.low_vram:
weight_streaming = getattr(obj, 'weight_streaming', False)
weight_streaming_budget_percentage = getattr(obj, 'weight_streaming_budget_percentage', None)
self.engine[model_name].load(weight_streaming, weight_streaming_budget_percentage)
if model_config['do_engine_refit'] and self.lora_loader:
# For low_vram, using on-demand load and unload for refit.
if self.low_vram:
assert self.engine[model_name].engine is None
self.engine[model_name].load()
self._refit_engine(obj, model_name, model_config)
if self.low_vram:
self.engine[model_name].unload()
# Load PyTorch models if torch-inference mode is enabled
self._load_torch_models()
# Reclaim GPU memory from torch cache
torch.cuda.empty_cache()
def calculate_max_device_memory(self):
max_device_memory = 0
for model_name, engine in self.engine.items():
if self.low_vram:
engine.load()
max_device_memory = max(max_device_memory, engine.engine.device_memory_size_v2)
if self.low_vram:
engine.unload()
return max_device_memory
def get_device_memory_sizes(self):
device_memory_sizes = {}
for model_name, engine in self.engine.items():
engine.load()
device_memory_sizes[model_name] = engine.engine.device_memory_size_v2
engine.unload()
return device_memory_sizes
def activate_engines(self, shared_device_memory=None):
if shared_device_memory is None:
max_device_memory = self.calculate_max_device_memory()
_, shared_device_memory = cudart.cudaMalloc(max_device_memory)
self.shared_device_memory = shared_device_memory
# Load and activate TensorRT engines
if not self.low_vram:
for engine in self.engine.values():
engine.activate(device_memory=self.shared_device_memory)
def run_engine(self, model_name, feed_dict):
engine = self.engine[model_name]
# CUDA graphs should be disabled when low_vram is enabled.
if self.low_vram:
assert self.use_cuda_graph == False
return engine.infer(feed_dict, self.stream, use_cuda_graph=self.use_cuda_graph)
def teardown(self):
for e in self.events.values():
cudart.cudaEventDestroy(e[0])
cudart.cudaEventDestroy(e[1])
for engine in self.engine.values():
engine.deallocate_buffers()
engine.deactivate()
engine.unload(verbose=False)
del engine
if self.shared_device_memory:
cudart.cudaFree(self.shared_device_memory)
for torch_model in self.torch_models.values():
torch_model.to("cpu")
del torch_model
cudart.cudaStreamDestroy(self.stream)
del self.stream
gc.collect()
torch.cuda.empty_cache()
def initialize_latents(self, batch_size, unet_channels, latent_height, latent_width, latents_dtype=torch.float32):
latents_shape = (batch_size, unet_channels, latent_height, latent_width)
latents = torch.randn(latents_shape, device=self.device, dtype=latents_dtype, generator=self.generator)
# Scale the initial noise by the standard deviation required by the scheduler
latents = latents * self.scheduler.init_noise_sigma
return latents
def save_image(self, images, pipeline, prompt, seed):
# Save image
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'
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."""
raise NotImplementedError("Please Implement the run method")
demo/Diffusion/demo_diffusion/pipeline/flux_pipeline.py
+1392
View File
File diff suppressed because it is too large Load Diff
demo/Diffusion/demo_diffusion/pipeline/model_memory_manager.py
+74
View File
@@ -0,0 +1,74 @@
#
# 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 cuda.bindings import runtime as cudart
class ModelMemoryManager:
"""
Context manager for efficiently loading and unloading models to optimize VRAM usage.
This class provides a context to temporarily load models into GPU memory for inference
and automatically unload them afterward. It's especially useful in low VRAM environments
where models need to be swapped in and out of GPU memory.
Args:
parent: The parent class instance that contains the model references and resources.
model_names (list): List of model names to load and unload.
low_vram (bool, optional): If True, enables VRAM optimization. If False, the context manager does nothing. Defaults to False.
"""
def __init__(self, parent, model_names, low_vram=False):
self.parent = parent
self.model_names = model_names
self.low_vram = low_vram
self.device_memories = {}
def __enter__(self):
if not self.low_vram:
return
for model_name in self.model_names:
if not self.parent.torch_fallback[model_name]:
# creating engine object (load from plan file)
self.parent.engine[model_name].load()
# allocate device memory
_, shared_device_memory = cudart.cudaMalloc(self.parent.device_memory_sizes[model_name])
self.device_memories[model_name] = shared_device_memory
# creating context
self.parent.engine[model_name].activate(device_memory=shared_device_memory)
# creating input and output buffer
self.parent.engine[model_name].allocate_buffers(
shape_dict=self.parent.shape_dicts[model_name], device=self.parent.device
)
else:
print(f"[I] Reloading torch model {model_name} from cpu.")
self.parent.torch_models[model_name] = self.parent.torch_models[model_name].to("cuda")
def __exit__(self, exc_type, exc_val, exc_tb):
if not self.low_vram:
return
for model_name in self.model_names:
if not self.parent.torch_fallback[model_name]:
self.parent.engine[model_name].deallocate_buffers()
self.parent.engine[model_name].deactivate()
self.parent.engine[model_name].unload()
cudart.cudaFree(self.device_memories.pop(model_name))
else:
print(f"[I] Offloading torch model {model_name} to cpu.")
self.parent.torch_models[model_name] = self.parent.torch_models[model_name].to("cpu")
torch.cuda.empty_cache()
demo/Diffusion/demo_diffusion/pipeline/stable_cascade_pipeline.py
+343
View File
@@ -0,0 +1,343 @@
#
# 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 inspect
import time
import tensorrt as trt
import torch
from cuda.bindings import runtime as cudart
from diffusers import DDPMWuerstchenScheduler
from demo_diffusion.model import (
CLIPWithProjModel,
UNetCascadeModel,
VQGANModel,
make_tokenizer,
)
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):
"""
Application showcasing the acceleration of Stable Cascade pipelines using NVidia TensorRT.
"""
def __init__(
self,
version='cascade',
pipeline_type=PIPELINE_TYPE.CASCADE_PRIOR,
latent_dim_scale=10.67,
lite=False,
**kwargs
):
"""
Initializes the Stable Cascade pipeline.
Args:
version (str):
The version of the pipeline. Should be one of [cascade]
pipeline_type (PIPELINE_TYPE):
Type of current pipeline.
latent_dim_scale (float):
Multiplier to determine the VQ latent space size from the image embeddings. If the image embeddings are
height=24 and width=24, the VQ latent shape needs to be height=int(24*10.67)=256 and
width=int(24*10.67)=256 in order to match the training conditions.
lite (bool):
Boolean indicating if the Lite Version of the Stage B and Stage C models is to be used
"""
super().__init__(
version=version,
pipeline_type=pipeline_type,
**kwargs
)
self.config['clip_hidden_states'] = True
# from Diffusers: https://github.com/huggingface/diffusers/blob/main/src/diffusers/pipelines/stable_cascade/pipeline_stable_cascade.py#L91C9-L91C41
self.latent_dim_scale = latent_dim_scale
self.lite = lite
def initializeModels(self, framework_model_dir, int8, fp8):
# Load text tokenizer(s)
self.tokenizer = make_tokenizer(self.version, self.pipeline_type, self.hf_token, framework_model_dir)
# Load pipeline models
models_args = {'version': self.version, 'pipeline': self.pipeline_type, 'device': self.device,
'hf_token': self.hf_token, 'verbose': self.verbose, 'framework_model_dir': framework_model_dir,
'max_batch_size': self.max_batch_size}
self.fp16 = False # TODO: enable FP16 mode for decoder model (requires strongly typed engine)
self.bf16 = True
if 'clip' in self.stages:
self.models['clip'] = CLIPWithProjModel(**models_args, fp16=self.fp16, bf16=self.bf16, output_hidden_states=self.config.get('clip_hidden_states', False), subfolder='text_encoder')
if 'unet' in self.stages:
self.models['unet'] = UNetCascadeModel(**models_args, fp16=self.fp16, bf16=self.bf16, lite=self.lite, do_classifier_free_guidance=self.do_classifier_free_guidance)
if 'vqgan' in self.stages:
self.models['vqgan'] = VQGANModel(**models_args, fp16=self.fp16, bf16=self.bf16, latent_dim_scale = self.latent_dim_scale)
def encode_prompt(self, prompt, negative_prompt, encoder='clip', pooled_outputs=False, output_hidden_states=False):
self.profile_start(encoder, color='green')
tokenizer = self.tokenizer
def tokenize(prompt, output_hidden_states):
text_inputs = tokenizer(
prompt,
padding="max_length",
max_length=tokenizer.model_max_length,
truncation=True,
return_tensors="pt",
)
text_input_ids = text_inputs.input_ids.type(torch.int32).to(self.device)
attention_mask = text_inputs.attention_mask.type(torch.int32).to(self.device)
text_hidden_states = None
if self.torch_inference:
outputs = self.torch_models[encoder](text_input_ids, attention_mask=attention_mask, output_hidden_states=output_hidden_states)
text_embeddings = outputs[0].clone()
if output_hidden_states:
hidden_state_layer = -1
text_hidden_states = outputs['hidden_states'][hidden_state_layer].clone()
else:
# NOTE: output tensor for CLIP must be cloned because it will be overwritten when called again for negative prompt
outputs = self.runEngine(encoder, {'input_ids': text_input_ids, 'attention_mask': attention_mask})
text_embeddings = outputs['text_embeddings'].clone()
if output_hidden_states:
text_hidden_states = outputs['hidden_states'].clone()
return text_embeddings, text_hidden_states
# Tokenize prompt
text_embeddings, text_hidden_states = tokenize(prompt, output_hidden_states)
if self.do_classifier_free_guidance:
# Tokenize negative prompt
uncond_embeddings, uncond_hidden_states = tokenize(negative_prompt, output_hidden_states)
# Concatenate the unconditional and text embeddings into a single batch to avoid doing two forward passes for classifier free guidance
text_embeddings = torch.cat([text_embeddings, uncond_embeddings])
if pooled_outputs:
pooled_output = text_embeddings
if output_hidden_states:
text_embeddings = torch.cat([text_hidden_states, uncond_hidden_states]) if self.do_classifier_free_guidance else text_hidden_states
self.profile_stop(encoder)
if pooled_outputs:
return text_embeddings, pooled_output
return text_embeddings
def denoise_latent(self,
latents,
pooled_embeddings,
text_embeddings=None,
image_embeds=None,
effnet=None,
denoiser='unet',
timesteps=None,
):
do_autocast = False
with torch.autocast('cuda', enabled=do_autocast):
self.profile_start(denoiser, color='blue')
for step_index, timestep in enumerate(timesteps):
# ratio input required for stable cascade prior
timestep_ratio = timestep.expand(latents.size(0)).to(latents.dtype)
# Expand the latents and timestep_ratio if we are doing classifier free guidance
latent_model_input = torch.cat([latents] * 2) if self.do_classifier_free_guidance else latents
timestep_ratio_input = torch.cat([timestep_ratio] * 2) if self.do_classifier_free_guidance else timestep_ratio
params = {"sample": latent_model_input, "timestep_ratio": timestep_ratio_input, "clip_text_pooled": pooled_embeddings}
if text_embeddings is not None:
params.update({'clip_text': text_embeddings})
if image_embeds is not None:
params.update({'clip_img': image_embeds})
if effnet is not None:
params.update({'effnet': effnet})
# Predict the noise residual
if self.torch_inference:
noise_pred = self.torch_models[denoiser](**params)['sample']
else:
noise_pred = self.runEngine(denoiser, params)['latent']
# Perform guidance
if self.do_classifier_free_guidance:
noise_pred_text, noise_pred_uncond = noise_pred.chunk(2)
noise_pred = noise_pred_uncond + self.guidance_scale * (noise_pred_text - noise_pred_uncond)
# from diffusers (prepare_extra_step_kwargs)
extra_step_kwargs = {}
if "eta" in set(inspect.signature(self.scheduler.step).parameters.keys()):
# TODO: configurable eta
eta = 0.0
extra_step_kwargs["eta"] = eta
if "generator" in set(inspect.signature(self.scheduler.step).parameters.keys()):
extra_step_kwargs["generator"] = self.generator
latents = self.scheduler.step(noise_pred, timestep_ratio, latents, **extra_step_kwargs, return_dict=False)[0]
latents = latents.to(dtype=torch.bfloat16 if self.bf16 else torch.float32)
self.profile_stop(denoiser)
return latents
def decode_latent(self, latents, model_name='vqgan'):
self.profile_start(model_name, color='red')
latents = self.models[model_name].scale_factor * latents
if self.torch_inference:
images = self.torch_models[model_name](latents)['sample']
else:
images = self.runEngine(model_name, {'latent': latents})['images']
self.profile_stop(model_name)
return images
def print_summary(self, denoising_steps, walltime_ms, batch_size):
print('|-----------------|--------------|')
print('| {:^15} | {:^12} |'.format('Module', 'Latency'))
print('|-----------------|--------------|')
for stage in self.stages:
stage_name = stage + ' x ' + str(denoising_steps) if stage == 'unet' else stage
print(
"| {:^15} | {:>9.2f} ms |".format(
stage_name, cudart.cudaEventElapsedTime(self.events[stage][0], self.events[stage][1])[1],
)
)
print('|-----------------|--------------|')
print('| {:^15} | {:>9.2f} ms |'.format('Pipeline', walltime_ms))
print('|-----------------|--------------|')
print('Throughput: {:.5f} image/s'.format(batch_size*1000./walltime_ms))
def infer(
self,
prompt,
negative_prompt,
image_height,
image_width,
image_embeddings=None,
warmup=False,
verbose=False,
save_image=True,
):
"""
Run the diffusion pipeline.
Args:
prompt (str):
The text prompt to guide image generation.
negative_prompt (str):
The prompt not to guide the image generation.
image_height (int):
Height (in pixels) of the image to be generated. Must be a multiple of 8.
image_width (int):
Width (in pixels) of the image to be generated. Must be a multiple of 8.
image_embeddings (`torch.FloatTensor` or `List[torch.FloatTensor]`):
Image Embeddings either extracted from an image or generated by a Prior Model.
warmup (bool):
Indicate if this is a warmup run.
verbose (bool):
Verbose in logging
save_image (bool):
Save the generated image (if applicable)
"""
if self.pipeline_type.is_cascade_decoder():
assert image_embeddings is not None, "Image Embeddings are required to run the decoder. Provided None"
assert len(prompt) == len(negative_prompt)
batch_size = len(prompt)
# Spatial dimensions of latent tensor
latent_height = image_height // 42
latent_width = image_width // 42
if image_embeddings is not None:
assert latent_height == image_embeddings.shape[-2]
assert latent_width == image_embeddings.shape[-1]
if self.generator and self.seed:
self.generator.manual_seed(self.seed)
num_inference_steps = self.denoising_steps
with torch.inference_mode(), trt.Runtime(TRT_LOGGER):
torch.cuda.synchronize()
e2e_tic = time.perf_counter()
denoise_kwargs = {}
# TODO: support custom timesteps
timesteps = None
if timesteps is not None:
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."
)
self.scheduler.set_timesteps(timesteps=timesteps, device=self.device)
assert self.denoising_steps == len(self.scheduler.timesteps)
else:
self.scheduler.set_timesteps(self.denoising_steps, device=self.device)
timesteps = self.scheduler.timesteps.to(self.device)
if isinstance(self.scheduler, DDPMWuerstchenScheduler):
timesteps = timesteps[:-1]
denoise_kwargs.update({'timesteps': timesteps})
# Initialize latents
latents_dtpye = torch.float16 if self.fp16 else torch.bfloat16 if self.bf16 else torch.float32
latents = self.initialize_latents(
batch_size=batch_size,
unet_channels=16 if self.pipeline_type.is_cascade_prior() else 4, # TODO: can we query "in_channels" from config
latent_height=latent_height if self.pipeline_type.is_cascade_prior() else int(latent_height * self.latent_dim_scale),
latent_width=latent_width if self.pipeline_type.is_cascade_prior() else int(latent_width * self.latent_dim_scale),
latents_dtype=latents_dtpye
)
# CLIP text encoder(s)
text_embeddings, pooled_embeddings = self.encode_prompt(prompt, negative_prompt,
encoder='clip', pooled_outputs=True, output_hidden_states=True)
if self.pipeline_type.is_cascade_prior():
denoise_kwargs.update({'text_embeddings': text_embeddings})
# image embeds
image_embeds_pooled = torch.zeros(batch_size, 1, 768, device=self.device, dtype=latents_dtpye)
image_embeds = (torch.cat([image_embeds_pooled, torch.zeros_like(image_embeds_pooled)]) if self.do_classifier_free_guidance else image_embeddings)
denoise_kwargs.update({'image_embeds': image_embeds})
else:
effnet = (torch.cat([image_embeddings, torch.zeros_like(image_embeddings)]) if self.do_classifier_free_guidance else image_embeddings)
denoise_kwargs.update({'effnet': effnet})
# UNet denoiser
latents = self.denoise_latent(latents, pooled_embeddings.unsqueeze(1), denoiser='unet', **denoise_kwargs)
if not self.return_latents:
images = self.decode_latent(latents)
torch.cuda.synchronize()
e2e_toc = time.perf_counter()
walltime_ms = (e2e_toc - e2e_tic) * 1000.
if not warmup:
self.print_summary(num_inference_steps, walltime_ms, batch_size)
if not self.return_latents and save_image:
# post-process images
images = ((images) * 255).clamp(0, 255).detach().permute(0, 2, 3, 1).round().type(torch.uint8).cpu().numpy()
self.save_image(images, self.pipeline_type.name.lower(), prompt, self.seed)
return (latents, walltime_ms) if self.return_latents else (images, walltime_ms)
demo/Diffusion/demo_diffusion/pipeline/stable_diffusion_35_pipeline.py
+952
View File
@@ -0,0 +1,952 @@
#
# 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 inspect
import os
import time
import warnings
from typing import Any, List, Union
import tensorrt as trt
import torch
from cuda.bindings import runtime as cudart
from diffusers.image_processor import VaeImageProcessor
from huggingface_hub import snapshot_download
from transformers import PreTrainedTokenizerBase
from demo_diffusion import path as path_module
from demo_diffusion.model import (
CLIPWithProjModel,
SD3ControlNet,
SD3TransformerModel,
T5Model,
VAEEncoderModel,
VAEModel,
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)
class SD3CannyImageProcessor(VaeImageProcessor):
def __init__(self):
super().__init__(do_normalize=False)
def preprocess(self, image, **kwargs):
image = super().preprocess(image, **kwargs)
image = image * 255 * 0.5 + 0.5
return image
def postprocess(self, image, do_denormalize=True, **kwargs):
do_denormalize = [True] * image.shape[0]
image = super().postprocess(image, **kwargs, do_denormalize=do_denormalize)
return image
class StableDiffusion35Pipeline(DiffusionPipeline):
"""
Application showcasing the acceleration of Stable Diffusion 3.5 pipelines using Nvidia TensorRT.
"""
def __init__(
self,
version: str,
pipeline_type=PIPELINE_TYPE.TXT2IMG,
guidance_scale: float = 7.0,
max_sequence_length: int = 256,
controlnet=None,
**kwargs,
):
"""
Initializes the Stable Diffusion 3.5 pipeline.
Args:
version (str):
The version of the pipeline. Should be one of ['3.5-medium', '3.5-large']
pipeline_type (PIPELINE_TYPE):
Type of current pipeline.
guidance_scale (`float`, defaults to 7.0):
Guidance scale is enabled by setting as > 1.
Higher guidance scale encourages to generate images that are closely linked to the text prompt, usually at the expense of lower image quality.
max_sequence_length (`int`, defaults to 256):
Maximum sequence length to use with the `prompt`.
controlnet (str):
Which ControlNet to use.
"""
super().__init__(version=version, pipeline_type=pipeline_type, controlnet=controlnet, **kwargs)
self.fp16 = True if not self.bf16 else False
self.force_weakly_typed_t5 = False
self.config["clip_g_torch_fallback"] = True
self.config["clip_l_torch_fallback"] = True
self.config["clip_hidden_states"] = True
self.controlnet = controlnet
self.guidance_scale = guidance_scale
self.do_classifier_free_guidance = self.guidance_scale > 1
self.max_sequence_length = max_sequence_length
@classmethod
def FromArgs(cls, args: argparse.Namespace, pipeline_type: PIPELINE_TYPE) -> StableDiffusion35Pipeline:
"""Factory method to construct a `StableDiffusion35Pipeline` object from parsed arguments.
Overrides:
DiffusionPipeline.FromArgs
"""
MAX_BATCH_SIZE = 4
DEVICE = "cuda"
DO_RETURN_LATENTS = False
# Resolve all paths.
controlnet_type = args.controlnet_type if "controlnet_type" in args else None
dd_path = path_module.resolve_path(
cls.get_model_names(pipeline_type, controlnet_type),
args,
pipeline_type,
cls._get_pipeline_uid(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,
controlnet=controlnet_type,
bf16=args.bf16,
low_vram=args.low_vram,
torch_fallback=args.torch_fallback,
weight_streaming=args.ws,
max_batch_size=MAX_BATCH_SIZE,
denoising_steps=args.denoising_steps,
scheduler=args.scheduler,
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, controlnet_type: str = None) -> List[str]:
"""Return a list of model names used by this pipeline.
Overrides:
DiffusionPipeline.get_model_names
"""
if pipeline_type.is_controlnet():
assert controlnet_type, "ControlNet type must be specified for ControlNet pipelines"
return ["clip_l", "clip_g", "t5", "transformer", "vae", "vae_encoder", f"controlnet_{controlnet_type}"]
return ["clip_l", "clip_g", "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. This pipeline is running in FP16."
)
hf_download_path = "-".join([load.get_path(self.version, self.pipeline_type.name), "tensorrt"])
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 == "transformer":
if model_config["use_fp8"]:
dirname = os.path.join(model_name, "fp8")
elif self.bf16:
dirname = os.path.join(model_name, "bf16")
elif "controlnet" in model_name:
hf_download_path_cnet = hf_download_path.replace("large", "controlnets")
dirname = f"controlnet_{self.controlnet}"
if "blur" in model_name:
pass
elif model_config["use_fp8"]:
dirname = os.path.join(dirname, "fp8")
elif self.bf16:
dirname = os.path.join(dirname, "bf16")
elif model_name in self.stages:
dirname = model_name
else:
raise ValueError(f"{model_name} not found in {self.stages}")
dirname = os.path.join("ONNX", dirname)
snapshot_download(
repo_id=hf_download_path if "controlnet" not in model_name else hf_download_path_cnet,
allow_patterns=os.path.join(dirname, "*"),
local_dir=base_dir,
token=self.hf_token,
)
# Rename directory from ONNX/<model_name> to <model_name>
saved_dir = os.path.join(base_dir, dirname)
model_dir = os.path.dirname(model_path)
os.rename(saved_dir, model_dir)
def load_resources(
self,
image_height: int,
image_width: int,
batch_size: int,
seed: int,
):
super().load_resources(image_height, image_width, batch_size, seed)
def _initialize_models(self, framework_model_dir, int8, fp8, fp4):
# Load text tokenizer(s)
self.tokenizer = make_tokenizer(
self.version,
self.pipeline_type,
self.hf_token,
framework_model_dir,
)
self.tokenizer2 = make_tokenizer(
self.version,
self.pipeline_type,
self.hf_token,
framework_model_dir,
subfolder="tokenizer_2",
)
self.tokenizer3 = make_tokenizer(
self.version,
self.pipeline_type,
self.hf_token,
framework_model_dir,
subfolder="tokenizer_3",
tokenizer_type="t5",
)
# Load pipeline models
models_args = {
"version": self.version,
"pipeline": self.pipeline_type,
"device": self.device,
"hf_token": self.hf_token,
"verbose": self.verbose,
"framework_model_dir": framework_model_dir,
"max_batch_size": self.max_batch_size,
}
self.bf16 = True if int8 or fp8 or fp4 else self.bf16
self.fp16 = True if not self.bf16 else False
self.tf32 = True
self.fp8 = fp8
self.int8 = int8
self.fp4 = fp4
if "clip_l" in self.stages:
self.models["clip_l"] = CLIPWithProjModel(
**models_args,
fp16=self.fp16,
bf16=self.bf16,
subfolder="text_encoder",
output_hidden_states=self.config.get("clip_hidden_states", False),
)
if "clip_g" in self.stages:
self.models["clip_g"] = CLIPWithProjModel(
**models_args,
fp16=self.fp16,
bf16=self.bf16,
subfolder="text_encoder_2",
output_hidden_states=self.config.get("clip_hidden_states", False),
)
if "t5" in self.stages:
# Known accuracy issues with FP16
self.models["t5"] = T5Model(
**models_args,
fp16=self.fp16,
bf16=self.bf16,
tf32=self.tf32,
subfolder="text_encoder_3",
text_maxlen=self.max_sequence_length,
weight_streaming=self.weight_streaming,
weight_streaming_budget_percentage=self.text_encoder_weight_streaming_budget_percentage,
)
if "transformer" in self.stages:
self.models["transformer"] = SD3TransformerModel(
**models_args,
fp16=self.fp16,
tf32=self.tf32,
bf16=self.bf16,
fp8=self.fp8,
int8=self.int8,
fp4=self.fp4,
text_maxlen=self.models["t5"].text_maxlen + self.models["clip_g"].text_maxlen,
weight_streaming=self.weight_streaming,
do_classifier_free_guidance=self.do_classifier_free_guidance,
)
if f"controlnet_{self.controlnet}" in self.stages:
self.models[f"controlnet_{self.controlnet}"] = SD3ControlNet(
**models_args,
fp16=self.fp16,
tf32=self.tf32,
bf16=self.bf16,
do_classifier_free_guidance=self.do_classifier_free_guidance,
controlnet=self.controlnet,
)
if "vae" in self.stages:
self.models["vae"] = VAEModel(**models_args, fp16=self.fp16, tf32=self.tf32, bf16=self.bf16)
self.vae_scale_factor = (
2 ** (len(self.models["vae"].config["block_out_channels"]) - 1) if "vae" in self.models else 8
)
self.patch_size = (
self.models["transformer"].config["patch_size"]
if "transformer" in self.stages and self.models["transformer"] is not None
else 2
)
if "vae_encoder" in self.stages:
self.models["vae_encoder"] = VAEEncoderModel(**models_args, fp16=False, tf32=self.tf32, bf16=self.bf16, do_classifier_free_guidance=self.do_classifier_free_guidance)
self.vae_latent_channels = (
self.models["vae"].config["latent_channels"]
if "vae" in self.stages and self.models["vae"] is not None
else 16
)
if self.controlnet and "canny" in self.controlnet:
self.image_processor = SD3CannyImageProcessor()
else:
self.image_processor = VaeImageProcessor(vae_scale_factor=self.vae_scale_factor)
def print_summary(self, denoising_steps, walltime_ms):
print("|-----------------|--------------|")
print("| {:^15} | {:^12} |".format("Module", "Latency"))
print("|-----------------|--------------|")
for stage in self.stages:
# controlnet is profiled in the denoising step
if "controlnet" in stage:
continue
stage_name = stage
if "transformer" in stage:
if f"controlnet_{self.controlnet}" in self.stages:
stage_name += '+cnet'
stage_name += ' x ' + str(denoising_steps)
print(
"| {:^15} | {:>9.2f} ms |".format(
stage_name, cudart.cudaEventElapsedTime(self.events[stage][0], self.events[stage][1])[1],
)
)
print("|-----------------|--------------|")
print("| {:^15} | {:>9.2f} ms |".format("Pipeline", walltime_ms))
print("|-----------------|--------------|")
print("Throughput: {:.5f} image/s".format(self.batch_size * 1000.0 / walltime_ms))
@staticmethod
def _tokenize(
tokenizer: PreTrainedTokenizerBase,
prompt: list[str],
max_sequence_length: int,
device: torch.device,
):
text_input_ids = tokenizer(
prompt,
padding="max_length",
max_length=max_sequence_length,
truncation=True,
add_special_tokens=True,
return_tensors="pt",
).input_ids
text_input_ids = text_input_ids.type(torch.int32)
untruncated_ids = tokenizer(
prompt,
padding="longest",
return_tensors="pt",
).input_ids.type(torch.int32)
if untruncated_ids.shape[-1] >= text_input_ids.shape[-1] and not torch.equal(text_input_ids, untruncated_ids):
removed_text = tokenizer.batch_decode(untruncated_ids[:, max_sequence_length - 1 : -1])
warnings.warn(
"The following part of your input was truncated because `max_sequence_length` is set to "
f" {max_sequence_length} tokens: {removed_text}"
)
text_input_ids = text_input_ids.to(device)
return text_input_ids
def _get_prompt_embed(
self,
prompt: list[str],
encoder_name: str,
domain="positive_prompt",
):
if encoder_name == "clip_l":
tokenizer = self.tokenizer
max_sequence_length = tokenizer.model_max_length
output_hidden_states = True
elif encoder_name == "clip_g":
tokenizer = self.tokenizer2
max_sequence_length = tokenizer.model_max_length
output_hidden_states = True
elif encoder_name == "t5":
tokenizer = self.tokenizer3
max_sequence_length = self.max_sequence_length
output_hidden_states = False
else:
raise NotImplementedError(f"encoder not found: {encoder_name}")
self.profile_start(encoder_name, color="green", domain=domain)
text_input_ids = self._tokenize(
tokenizer=tokenizer,
prompt=prompt,
device=self.device,
max_sequence_length=max_sequence_length,
)
text_hidden_states = None
if self.torch_inference or self.torch_fallback[encoder_name]:
outputs = self.torch_models[encoder_name](
text_input_ids,
output_hidden_states=output_hidden_states,
)
text_embeddings = outputs[0].clone()
if output_hidden_states:
text_hidden_states = outputs["hidden_states"][-2].clone()
else:
# NOTE: output tensor for the encoder must be cloned because it will be overwritten when called again for prompt2
outputs = self.run_engine(encoder_name, {"input_ids": text_input_ids})
text_embeddings = outputs["text_embeddings"].clone()
if output_hidden_states:
text_hidden_states = outputs["hidden_states"].clone()
self.profile_stop(encoder_name)
return text_hidden_states, text_embeddings
@staticmethod
def _duplicate_text_embed(
prompt_embed: torch.Tensor,
batch_size: int,
num_images_per_prompt: int,
pooled_prompt_embed: torch.Tensor | None = None,
):
_, seq_len, _ = prompt_embed.shape
# duplicate text embeddings for each generation per prompt, using mps friendly method
prompt_embed = prompt_embed.repeat(1, num_images_per_prompt, 1)
prompt_embed = prompt_embed.view(batch_size * num_images_per_prompt, seq_len, -1)
if pooled_prompt_embed is not None:
pooled_prompt_embed = pooled_prompt_embed.repeat(1, num_images_per_prompt, 1)
pooled_prompt_embed = pooled_prompt_embed.view(batch_size * num_images_per_prompt, -1)
return prompt_embed, pooled_prompt_embed
def encode_prompt(
self,
prompt: list[str],
negative_prompt: list[str] | None = None,
num_images_per_prompt: int = 1,
):
clip_l_prompt_embed, clip_l_pooled_embed = self._get_prompt_embed(
prompt=prompt,
encoder_name="clip_l",
)
prompt_embed, pooled_prompt_embed = self._duplicate_text_embed(
prompt_embed=clip_l_prompt_embed.clone(),
pooled_prompt_embed=clip_l_pooled_embed.clone(),
num_images_per_prompt=num_images_per_prompt,
batch_size=self.batch_size,
)
clip_g_prompt_embed, clip_g_pooled_embed = self._get_prompt_embed(
prompt=prompt,
encoder_name="clip_g",
)
prompt_2_embed, pooled_prompt_2_embed = self._duplicate_text_embed(
prompt_embed=clip_g_prompt_embed.clone(),
pooled_prompt_embed=clip_g_pooled_embed.clone(),
batch_size=self.batch_size,
num_images_per_prompt=num_images_per_prompt,
)
_, t5_prompt_embed = self._get_prompt_embed(
prompt=prompt,
encoder_name="t5",
)
t5_prompt_embed, _ = self._duplicate_text_embed(
prompt_embed=t5_prompt_embed.clone(),
batch_size=self.batch_size,
num_images_per_prompt=num_images_per_prompt,
)
clip_prompt_embeds = torch.cat([prompt_embed, prompt_2_embed], dim=-1)
clip_prompt_embeds = torch.nn.functional.pad(
clip_prompt_embeds, (0, t5_prompt_embed.shape[-1] - clip_prompt_embeds.shape[-1])
)
prompt_embeds = torch.cat([clip_prompt_embeds, t5_prompt_embed], dim=-2)
pooled_prompt_embeds = torch.cat([pooled_prompt_embed, pooled_prompt_2_embed], dim=-1)
if negative_prompt is None:
negative_prompt = ""
clip_l_negative_prompt_embed, clip_l_negative_pooled_embed = self._get_prompt_embed(
prompt=negative_prompt,
encoder_name="clip_l",
)
negative_prompt_embed, negative_pooled_prompt_embed = self._duplicate_text_embed(
prompt_embed=clip_l_negative_prompt_embed.clone(),
pooled_prompt_embed=clip_l_negative_pooled_embed.clone(),
batch_size=self.batch_size,
num_images_per_prompt=num_images_per_prompt,
)
clip_g_negative_prompt_embed, clip_g_negative_pooled_embed = self._get_prompt_embed(
prompt=negative_prompt,
encoder_name="clip_g",
)
negative_prompt_2_embed, negative_pooled_prompt_2_embed = self._duplicate_text_embed(
prompt_embed=clip_g_negative_prompt_embed.clone(),
pooled_prompt_embed=clip_g_negative_pooled_embed.clone(),
batch_size=self.batch_size,
num_images_per_prompt=num_images_per_prompt,
)
_, t5_negative_prompt_embed = self._get_prompt_embed(
prompt=negative_prompt,
encoder_name="t5",
)
t5_negative_prompt_embed, _ = self._duplicate_text_embed(
prompt_embed=t5_negative_prompt_embed.clone(),
batch_size=self.batch_size,
num_images_per_prompt=num_images_per_prompt,
)
negative_clip_prompt_embeds = torch.cat([negative_prompt_embed, negative_prompt_2_embed], dim=-1)
negative_clip_prompt_embeds = torch.nn.functional.pad(
negative_clip_prompt_embeds,
(0, t5_negative_prompt_embed.shape[-1] - negative_clip_prompt_embeds.shape[-1]),
)
negative_prompt_embeds = torch.cat([negative_clip_prompt_embeds, t5_negative_prompt_embed], dim=-2)
negative_pooled_prompt_embeds = torch.cat(
[negative_pooled_prompt_embed, negative_pooled_prompt_2_embed], dim=-1
)
return prompt_embeds, negative_prompt_embeds, pooled_prompt_embeds, negative_pooled_prompt_embeds
@staticmethod
def initialize_latents(
batch_size: int,
num_channels_latents: int,
latent_height: int,
latent_width: int,
device: torch.device,
generator: torch.Generator,
dtype=torch.float32,
layout=torch.strided,
) -> torch.Tensor:
latents_shape = (batch_size, num_channels_latents, latent_height, latent_width)
latents = torch.randn(
latents_shape,
dtype=dtype,
device="cuda",
generator=generator,
layout=layout,
).to(device)
return latents
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps
@staticmethod
def retrieve_timesteps(
scheduler,
num_inference_steps: int | None = None,
device: str | torch.device | None = None,
timesteps: list[int] | None = None,
sigmas: list[float] | None = None,
**kwargs,
):
r"""
Calls the scheduler's `set_timesteps` method and retrieves timesteps from the scheduler after the call. Handles
custom timesteps. Any kwargs will be supplied to `scheduler.set_timesteps`.
Args:
scheduler (`SchedulerMixin`):
The scheduler to get timesteps from.
num_inference_steps (`int`):
The number of diffusion steps used when generating samples with a pre-trained model. If used, `timesteps`
must be `None`.
device (`str` or `torch.device`, *optional*):
The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
timesteps (`List[int]`, *optional*):
Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed,
`num_inference_steps` and `sigmas` must be `None`.
sigmas (`List[float]`, *optional*):
Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed,
`num_inference_steps` and `timesteps` must be `None`.
Returns:
`Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
second element is the number of inference steps.
"""
if timesteps is not None and sigmas is not None:
raise ValueError(
"Only one of `timesteps` or `sigmas` can be passed. Please choose one to set custom values"
)
if timesteps is not None:
accepts_timesteps = "timesteps" in set(inspect.signature(scheduler.set_timesteps).parameters.keys())
if not accepts_timesteps:
raise ValueError(
f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom"
f" timestep schedules. Please check whether you are using the correct scheduler."
)
scheduler.set_timesteps(timesteps=timesteps, device=device, **kwargs)
timesteps = scheduler.timesteps
num_inference_steps = len(timesteps)
elif sigmas is not None:
accept_sigmas = "sigmas" in set(inspect.signature(scheduler.set_timesteps).parameters.keys())
if not accept_sigmas:
raise ValueError(
f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom"
f" sigmas schedules. Please check whether you are using the correct scheduler."
)
scheduler.set_timesteps(sigmas=sigmas, device=device, **kwargs)
timesteps = scheduler.timesteps
num_inference_steps = len(timesteps)
else:
scheduler.set_timesteps(num_inference_steps, device=device, **kwargs)
timesteps = scheduler.timesteps
return timesteps, num_inference_steps
def get_control_block_samples(self, params_controlnet, controlnet_name="controlnet"):
# Predict the controlnet block samples
if self.torch_inference or self.torch_fallback[controlnet_name]:
block_samples = self.torch_models[controlnet_name](**params_controlnet)
else:
block_samples = self.run_engine(controlnet_name, params_controlnet)["controlnet_block_samples"].clone()
return block_samples
def denoise_latents(
self,
latents: torch.Tensor,
prompt_embeds: torch.Tensor,
pooled_prompt_embeds: torch.Tensor,
timesteps: torch.FloatTensor,
guidance_scale: float,
denoiser="transformer",
control_image=None,
controlnet_scale=None,
controlnet_keep=None,
) -> torch.Tensor:
do_autocast = self.torch_inference != "" and self.models[denoiser].fp16
with torch.autocast("cuda", enabled=do_autocast):
self.profile_start(denoiser, color="blue")
for step_index, timestep in enumerate(timesteps):
# expand the latents as we are doing classifier free guidance
latents_model_input = torch.cat([latents] * 2) if self.do_classifier_free_guidance else latents
# broadcast to batch dimension in a way that's compatible with ONNX/Core ML
timestep_inp = timestep.expand(latents_model_input.shape[0])
controlnet_name = f"controlnet_{self.controlnet}"
if control_image is not None:
cond_scale = controlnet_scale * controlnet_keep[step_index]
cast_to = (
torch.float16
if self.models[controlnet_name].fp16
else torch.bfloat16 if self.models[controlnet_name].bf16 else torch.float32
)
params_controlnet = {
"hidden_states": latents_model_input,
"timestep": timestep_inp,
"pooled_projections": pooled_prompt_embeds,
"controlnet_cond": control_image,
"conditioning_scale": cond_scale.to(self.device).to(cast_to),
}
control_block_samples = self.get_control_block_samples(params_controlnet, controlnet_name)
else:
latent_shape = latents_model_input.shape
# Initialize control block samples with zeros. Hard-coding some dimensions that can only be queried if a controlnet is used.
control_block_samples = torch.zeros(
self.models["transformer"].num_controlnet_layers,
latent_shape[0], # batch size
latent_shape[2] // 2 * latent_shape[3] // 2,
self.models["transformer"].config["num_attention_heads"]
* self.models["transformer"].config["attention_head_dim"],
dtype=latents.dtype,
device=latents.device,
)
params = {
"hidden_states": latents_model_input,
"timestep": timestep_inp,
"encoder_hidden_states": prompt_embeds,
"pooled_projections": pooled_prompt_embeds,
"block_controlnet_hidden_states": control_block_samples,
}
# Predict the noise residual
if self.torch_inference or self.torch_fallback[denoiser]:
noise_pred = self.torch_models[denoiser](**params)["sample"]
else:
noise_pred = self.run_engine(denoiser, params)["latent"]
# perform guidance
if self.do_classifier_free_guidance:
noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)
# compute the previous noisy sample x_t -> x_t-1
latents = self.scheduler.step(noise_pred, timestep, latents, return_dict=False)[0]
self.profile_stop(denoiser)
return latents
def prepare_image(
self,
image,
width,
height,
device,
dtype,
):
image = self.image_processor.preprocess(image, height=height, width=width)
image = image.to(device=device, dtype=dtype)
if self.do_classifier_free_guidance:
image = torch.cat([image] * 2)
return image
def encode_image(self, input_image, model_name="vae_encoder"):
self.profile_start(model_name, color="red")
cast_to = (
torch.float16
if self.models[model_name].fp16
else torch.bfloat16 if self.models[model_name].bf16 else torch.float32
)
input_image = input_image.to(dtype=cast_to)
if self.torch_inference or self.torch_fallback[model_name]:
image_latents = self.torch_models[model_name](input_image)
else:
image_latents = self.run_engine(model_name, {"images": input_image})["latent"]
image_latents = (image_latents - self.models["vae"].config["shift_factor"]) * self.models[
"vae"
].config["scaling_factor"]
self.profile_stop(model_name)
return image_latents
def decode_latents(self, latents: torch.Tensor, decoder="vae") -> torch.Tensor:
cast_to = (
torch.float16
if self.models[decoder].fp16
else torch.bfloat16
if self.models[decoder].bf16
else torch.float32
)
latents = latents.to(dtype=cast_to)
self.profile_start(decoder, color="red")
if self.torch_inference or self.torch_fallback[decoder]:
images = self.torch_models[decoder](latents, return_dict=False)[0]
else:
images = self.run_engine(decoder, {"latent": latents})["images"]
self.profile_stop(decoder)
return images
def infer(
self,
prompt: list[str],
negative_prompt: list[str],
image_height: int,
image_width: int,
control_image=None,
controlnet_scale=None,
control_guidance_start: Union[float, List[float]] = 0.0,
control_guidance_end: Union[float, List[float]] = 1.0,
warmup=False,
save_image=True,
):
"""
Run the diffusion pipeline.
Args:
prompt (list[str]):
The text prompt to guide image generation.
negative_prompt (list[str]):
The prompt not to guide the image generation.
image_height (int):
Height (in pixels) of the image to be generated. Must be a multiple of 8.
image_width (int):
Width (in pixels) of the image to be generated. Must be a multiple of 8.
control_image (PIL.Image.Image):
The control image to guide the image generation.
controlnet_scale (torch.Tensor):
A tensor which contains ControlNet scale, essential for multi ControlNet.
control_guidance_start (`float` or `List[float]`, *optional*, defaults to 0.0):
The percentage of total steps at which the ControlNet starts applying.
control_guidance_end (`float` or `List[float]`, *optional*, defaults to 1.0):
The percentage of total steps at which the ControlNet stops applying.
warmup (bool):
Indicate if this is a warmup run.
save_image (bool):
Save the generated image (if applicable)
"""
assert len(prompt) == len(negative_prompt)
self.batch_size = len(prompt)
# Spatial dimensions of latent tensor
assert image_height % (self.vae_scale_factor * self.patch_size) == 0, (
f"image height not supported {image_height}"
)
assert image_width % (self.vae_scale_factor * self.patch_size) == 0, f"image width not supported {image_width}"
latent_height = int(image_height) // self.vae_scale_factor
latent_width = int(image_width) // self.vae_scale_factor
with torch.inference_mode(), trt.Runtime(TRT_LOGGER):
torch.cuda.synchronize()
e2e_tic = time.perf_counter()
# 1. encode inputs
with self.model_memory_manager(["clip_g", "clip_l", "t5"], low_vram=self.low_vram):
(
prompt_embeds,
negative_prompt_embeds,
pooled_prompt_embeds,
negative_pooled_prompt_embeds,
) = self.encode_prompt(
prompt=prompt,
negative_prompt=negative_prompt,
num_images_per_prompt=1,
)
# do classifier free guidance
if self.do_classifier_free_guidance:
prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds], dim=0)
pooled_prompt_embeds = torch.cat([negative_pooled_prompt_embeds, pooled_prompt_embeds], dim=0)
# 2. Prepare latent variables
num_channels_latents = self.models["transformer"].config["in_channels"]
latents = self.initialize_latents(
batch_size=self.batch_size,
num_channels_latents=num_channels_latents,
latent_height=latent_height,
latent_width=latent_width,
device=prompt_embeds.device,
generator=self.generator,
dtype=torch.float16 if self.fp16 else torch.bfloat16 if self.bf16 else torch.float32,
)
# 3. Prepare timesteps
timesteps, num_inference_steps = self.retrieve_timesteps(
scheduler=self.scheduler,
num_inference_steps=self.denoising_steps,
device=self.device,
sigmas=None,
)
# 4. Prepare control image
controlnet_keep = []
if control_image is not None:
# Process controlnet_scales
for i in range(len(timesteps)):
keeps = [
1.0 - float(i / len(timesteps) < s or (i + 1) / len(timesteps) > e)
for s, e in zip([control_guidance_start], [control_guidance_end])
]
controlnet_keep.append(keeps[0])
control_image = self.prepare_image(
image=control_image,
width=image_width,
height=image_height,
device=self.device,
dtype=torch.float16 if self.fp16 else torch.bfloat16 if self.bf16 else torch.float32,
)
with self.model_memory_manager(["vae_encoder"], low_vram=self.low_vram):
control_image = self.encode_image(control_image)
# 5. Denoise
denoiser_list = ["transformer", f"controlnet_{self.controlnet}"] if self.controlnet else ["transformer"]
with self.model_memory_manager(denoiser_list, low_vram=self.low_vram):
latents = self.denoise_latents(
latents=latents,
prompt_embeds=prompt_embeds,
pooled_prompt_embeds=pooled_prompt_embeds,
timesteps=timesteps,
guidance_scale=self.guidance_scale,
control_image=control_image,
# TODO: support multiple controlnets
controlnet_scale=controlnet_scale,
controlnet_keep=controlnet_keep,
)
# 6. Decode Latents
latents = (latents / self.models["vae"].config["scaling_factor"]) + self.models["vae"].config[
"shift_factor"
]
with self.model_memory_manager(["vae"], low_vram=self.low_vram):
images = self.decode_latents(latents)
torch.cuda.synchronize()
e2e_toc = time.perf_counter()
walltime_ms = (e2e_toc - e2e_tic) * 1000.0
if not warmup:
self.print_summary(
num_inference_steps,
walltime_ms,
)
if save_image:
# post-process images
images = (
((images + 1) * 255 / 2)
.clamp(0, 255)
.detach()
.permute(0, 2, 3, 1)
.round()
.type(torch.uint8)
.cpu()
.numpy()
)
self.save_image(images, self.pipeline_type.name.lower(), prompt, self.seed)
return images, walltime_ms
def run(
self,
prompt: list[str],
negative_prompt: list[str],
height: int,
width: int,
batch_count: int,
num_warmup_runs: int,
use_cuda_graph: bool,
**kwargs,
):
num_warmup_runs = max(1, num_warmup_runs) if use_cuda_graph else num_warmup_runs
if num_warmup_runs > 0:
print("[I] Warming up ..")
for _ in range(num_warmup_runs):
self.infer(prompt, negative_prompt, height, width, warmup=True, **kwargs)
for _ in range(batch_count):
print("[I] Running StableDiffusion 3.5 pipeline")
if self.nvtx_profile:
cudart.cudaProfilerStart()
self.infer(prompt, negative_prompt, height, width, warmup=False, **kwargs)
if self.nvtx_profile:
cudart.cudaProfilerStop()
demo/Diffusion/demo_diffusion/pipeline/stable_diffusion_3_pipeline.py
+599
View File
@@ -0,0 +1,599 @@
#
# 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 math
import os
import pathlib
import time
import nvtx
import tensorrt as trt
import torch
from cuda.bindings import runtime as cudart
import demo_diffusion.engine as engine_module
import demo_diffusion.image as image_module
from demo_diffusion.model import (
SD3_CLIPGModel,
SD3_CLIPLModel,
SD3_MMDiTModel,
SD3_T5XXLModel,
SD3_VAEDecoderModel,
SD3_VAEEncoderModel,
get_clip_embedding_dim,
)
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:
"""
Application showcasing the acceleration of Stable Diffusion 3 pipelines using NVidia TensorRT.
"""
def __init__(
self,
version='sd3',
pipeline_type=PIPELINE_TYPE.TXT2IMG,
max_batch_size=16,
shift=1.0,
cfg_scale=5,
denoising_steps=50,
denoising_percentage=0.6,
input_image=None,
device='cuda',
output_dir='.',
hf_token=None,
verbose=False,
nvtx_profile=False,
use_cuda_graph=False,
framework_model_dir='pytorch_model',
torch_inference='',
):
"""
Initializes the Stable Diffusion 3 pipeline.
Args:
version (str):
The version of the pipeline. Should be one of ['sd3]
pipeline_type (PIPELINE_TYPE):
Type of current pipeline.
max_batch_size (int):
Maximum batch size for dynamic batch engine.
shift (float):
Shift parameter for MMDiT model. Default: 1.0
cfg_scale (int):
CFG Scale used for denoising. Default: 5
denoising_steps (int):
Number of denoising steps. Default: 1.0
denoising_percentage (float):
Denoising percentage. Default: 0.6
input_image (float):
Input image for conditioning. Default: None
device (str):
PyTorch device to run inference. Default: 'cuda'
output_dir (str):
Output directory for log files and image artifacts
hf_token (str):
HuggingFace User Access Token to use for downloading Stable Diffusion model checkpoints.
verbose (bool):
Enable verbose logging.
nvtx_profile (bool):
Insert NVTX profiling markers.
use_cuda_graph (bool):
Use CUDA graph to capture engine execution and then launch inference
framework_model_dir (str):
cache directory for framework checkpoints
torch_inference (str):
Run inference with PyTorch (using specified compilation mode) instead of TensorRT.
"""
self.max_batch_size = max_batch_size
self.shift = shift
self.cfg_scale = cfg_scale
self.denoising_steps = denoising_steps
self.input_image = input_image
self.denoising_percentage = denoising_percentage if input_image is not None else 1.0
self.framework_model_dir = framework_model_dir
self.output_dir = output_dir
for directory in [self.framework_model_dir, self.output_dir]:
if not os.path.exists(directory):
print(f"[I] Create directory: {directory}")
pathlib.Path(directory).mkdir(parents=True)
self.hf_token = hf_token
self.device = device
self.verbose = verbose
self.nvtx_profile = nvtx_profile
self.version = version
# Pipeline type
self.pipeline_type = pipeline_type
self.stages = ['clip_g', 'clip_l', 't5xxl', 'transformer', 'vae_decoder']
if input_image is not None:
self.stages = ['vae_encoder'] + self.stages
self.config = {}
self.config['clip_hidden_states'] = True
self.config['t5xxl_torch_fallback'] = True
self.config['vae_encoder_torch_fallback'] = True
self.torch_inference = torch_inference
if self.torch_inference:
torch._inductor.config.conv_1x1_as_mm = True
torch._inductor.config.coordinate_descent_tuning = True
torch._inductor.config.epilogue_fusion = False
torch._inductor.config.coordinate_descent_check_all_directions = True
self.use_cuda_graph = use_cuda_graph
# initialized in loadEngines()
self.models = {}
self.torch_models = {}
self.engine = {}
self.shared_device_memory = None
# initialized in loadResources()
self.events = {}
self.generator = None
self.markers = {}
self.seed = None
self.stream = None
self.tokenizer = None
def loadResources(self, image_height, image_width, batch_size, seed):
# Initialize noise generator
if seed:
self.seed = seed
self.generator = torch.Generator(device="cuda").manual_seed(seed)
# Create CUDA events and stream
for stage in self.stages:
self.events[stage] = [cudart.cudaEventCreate()[1], cudart.cudaEventCreate()[1]]
self.stream = cudart.cudaStreamCreate()[1]
# Allocate TensorRT I/O buffers
if not self.torch_inference:
for model_name, obj in self.models.items():
if self.torch_fallback[model_name]:
continue
self.engine[model_name].allocate_buffers(shape_dict=obj.get_shape_dict(batch_size, image_height, image_width), device=self.device)
def teardown(self):
for e in self.events.values():
cudart.cudaEventDestroy(e[0])
cudart.cudaEventDestroy(e[1])
for engine in self.engine.values():
del engine
if self.shared_device_memory:
cudart.cudaFree(self.shared_device_memory)
cudart.cudaStreamDestroy(self.stream)
del self.stream
def getOnnxPath(self, model_name, onnx_dir, opt=True, suffix=''):
onnx_model_dir = os.path.join(onnx_dir, 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 getEnginePath(self, model_name, engine_dir, enable_refit=False, suffix=''):
return os.path.join(engine_dir, model_name+suffix+('.refit' if enable_refit else '')+'.trt'+trt.__version__+'.plan')
def loadEngines(
self,
engine_dir,
framework_model_dir,
onnx_dir,
onnx_opset,
opt_batch_size,
opt_image_height,
opt_image_width,
static_batch=False,
static_shape=True,
enable_all_tactics=False,
timing_cache=None,
**_kwargs,
):
"""
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):
Batch size to optimize for during engine building.
opt_image_height (int):
Image height to optimize for during engine building. Must be a multiple of 8.
opt_image_width (int):
Image width to optimize for during engine building. Must be a multiple of 8.
static_batch (bool):
Build engine only for specified opt_batch_size.
static_shape (bool):
Build engine only for specified opt_image_height & opt_image_width. Default = True.
enable_all_tactics (bool):
Enable all tactic sources during TensorRT engine builds.
timing_cache (str):
Path to the timing cache to speed up TensorRT build.
"""
# 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)
# Load pipeline models
models_args = {'version': self.version, 'pipeline': self.pipeline_type, 'device': self.device,
'hf_token': self.hf_token, 'verbose': self.verbose, 'framework_model_dir': framework_model_dir,
'max_batch_size': self.max_batch_size}
# Load text tokenizer
self.tokenizer = SD3Tokenizer()
# Load text encoders
if 'clip_g' in self.stages:
self.models['clip_g'] = SD3_CLIPGModel(**models_args, fp16=True, pooled_output=True)
if 'clip_l' in self.stages:
self.models['clip_l'] = SD3_CLIPLModel(**models_args, fp16=True, pooled_output=True)
if 't5xxl' in self.stages:
self.models['t5xxl'] = SD3_T5XXLModel(**models_args, fp16=True, embedding_dim=get_clip_embedding_dim(self.version, self.pipeline_type))
# Load Transformer model
if 'transformer' in self.stages:
self.models['transformer'] = SD3_MMDiTModel(**models_args, fp16=True, shift=self.shift)
# Load VAE Encoder model
if 'vae_encoder' in self.stages:
self.models['vae_encoder'] = SD3_VAEEncoderModel(**models_args, fp16=True)
# Load VAE Decoder model
if 'vae_decoder' in self.stages:
self.models['vae_decoder'] = SD3_VAEDecoderModel(**models_args, fp16=True)
# Configure pipeline models to load
model_names = self.models.keys()
# Torch fallback
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]))
onnx_path = dict(zip(model_names, [self.getOnnxPath(model_name, onnx_dir, opt=False) for model_name in model_names]))
onnx_opt_path = dict(zip(model_names, [self.getOnnxPath(model_name, onnx_dir) for model_name in model_names]))
engine_path = dict(zip(model_names, [self.getEnginePath(model_name, engine_dir) for model_name in model_names]))
for model_name, obj in self.models.items():
if self.torch_fallback[model_name]:
continue
# Export models to ONNX
do_export_onnx = not os.path.exists(engine_path[model_name]) and not os.path.exists(onnx_opt_path[model_name])
if do_export_onnx:
obj.export_onnx(onnx_path[model_name], onnx_opt_path[model_name], onnx_opset, opt_image_height, opt_image_width, static_shape=static_shape)
# Build TensorRT engines
for model_name, obj in self.models.items():
if self.torch_fallback[model_name]:
continue
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}
engine.build(onnx_opt_path[model_name],
input_profile=obj.get_input_profile(
opt_batch_size, opt_image_height, opt_image_width,
static_batch=static_batch, static_shape=static_shape
),
enable_all_tactics=enable_all_tactics,
timing_cache=timing_cache,
update_output_names=update_output_names,
verbose=self.verbose
)
self.engine[model_name] = engine
# Load TensorRT engines
for model_name, obj in self.models.items():
if self.torch_fallback[model_name]:
continue
self.engine[model_name].load()
# Load torch models
for model_name, obj in self.models.items():
if self.torch_fallback[model_name] or model_name == 'transformer':
self.torch_models[model_name] = obj.get_model(torch_inference=self.torch_inference)
def calculateMaxDeviceMemory(self):
max_device_memory = 0
for model_name, engine in self.engine.items():
max_device_memory = max(max_device_memory, engine.engine.device_memory_size_v2)
return max_device_memory
def activateEngines(self, shared_device_memory=None):
if shared_device_memory is None:
max_device_memory = self.calculateMaxDeviceMemory()
_, shared_device_memory = cudart.cudaMalloc(max_device_memory)
self.shared_device_memory = shared_device_memory
# Load and activate TensorRT engines
for engine in self.engine.values():
engine.activate(device_memory=self.shared_device_memory)
def runEngine(self, model_name, feed_dict):
engine = self.engine[model_name]
return engine.infer(feed_dict, self.stream, use_cuda_graph=self.use_cuda_graph)
def initialize_latents(self, batch_size, unet_channels, latent_height, latent_width):
return torch.ones(batch_size, unet_channels, latent_height, latent_width, device="cuda") * 0.0609
def profile_start(self, name, color='blue'):
if self.nvtx_profile:
self.markers[name] = nvtx.start_range(message=name, color=color)
if name in self.events:
cudart.cudaEventRecord(self.events[name][0], 0)
def profile_stop(self, name):
if name in self.events:
cudart.cudaEventRecord(self.events[name][1], 0)
if self.nvtx_profile:
nvtx.end_range(self.markers[name])
def print_summary(self, denoising_steps, walltime_ms, batch_size):
print('|-----------------|--------------|')
print('| {:^15} | {:^12} |'.format('Module', 'Latency'))
print('|-----------------|--------------|')
for stage in self.stages:
stage_name = stage
if "transformer" in stage:
stage_name += ' x ' + str(denoising_steps)
print(
"| {:^15} | {:>9.2f} ms |".format(
stage_name, cudart.cudaEventElapsedTime(self.events[stage][0], self.events[stage][1])[1],
)
)
print('|-----------------|--------------|')
print('| {:^15} | {:>9.2f} ms |'.format('Pipeline', walltime_ms))
print('|-----------------|--------------|')
print('Throughput: {:.5f} image/s'.format(batch_size*1000./walltime_ms))
def save_image(self, images, pipeline, prompt, seed):
# 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'
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):
self.profile_start(model_name, color='green')
tokens = list(map(lambda a: a[0], token_weight_pairs[0]))
tokens = torch.tensor([tokens], dtype=torch.int64, device=self.device)
if self.torch_inference or self.torch_fallback[model_name]:
out, pooled = self.torch_models[model_name](tokens)
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:
first_pooled = pooled[0:1].cuda()
else:
first_pooled = pooled
output = [out[0:1]]
return torch.cat(output, dim=-2).cuda(), first_pooled
def tokenize(prompt):
tokens = self.tokenizer.tokenize_with_weights(prompt)
l_out, l_pooled = encode_token_weights('clip_l', tokens["l"])
g_out, g_pooled = encode_token_weights('clip_g', tokens["g"])
t5_out, _ = encode_token_weights('t5xxl', tokens["t5xxl"])
lg_out = torch.cat([l_out, g_out], dim=-1)
lg_out = torch.nn.functional.pad(lg_out, (0, 4096 - lg_out.shape[-1]))
return torch.cat([lg_out, t5_out], dim=-2), torch.cat((l_pooled, g_pooled), dim=-1)
conditioning = tokenize(prompt[0])
neg_conditioning = tokenize(negative_prompt[0])
return conditioning, neg_conditioning
def denoise_latent(self, latent, conditioning, neg_conditioning, model_name='transformer'):
def get_noise(latent):
return torch.randn(latent.size(), dtype=torch.float32, layout=latent.layout, generator=self.generator, device="cuda").to(latent.dtype)
def get_sigmas(sampling, steps):
start = sampling.timestep(sampling.sigma_max)
end = sampling.timestep(sampling.sigma_min)
timesteps = torch.linspace(start, end, steps)
sigs = []
for x in range(len(timesteps)):
ts = timesteps[x]
sigs.append(sampling.sigma(ts))
sigs += [0.0]
return torch.FloatTensor(sigs)
def max_denoise(sigmas):
max_sigma = float(self.torch_models[model_name].model_sampling.sigma_max)
sigma = float(sigmas[0])
return math.isclose(max_sigma, sigma, rel_tol=1e-05) or sigma > max_sigma
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])
sigma = torch.cat([timestep, timestep])
c_crossattn = torch.cat([cond["c_crossattn"], uncond["c_crossattn"]])
y = torch.cat([cond["y"], uncond["y"]])
if self.torch_inference or self.torch_fallback[model_name]:
with torch.autocast("cuda", dtype=torch.float16):
batched = self.torch_models[model_name](sample, sigma, c_crossattn=c_crossattn, y=y)
else:
input_dict = {'sample': sample, 'sigma': sigma, 'c_crossattn': c_crossattn, 'y': y}
batched = self.runEngine(model_name, input_dict)['latent']
# Then split and apply CFG Scaling
pos_out, neg_out = batched.chunk(2)
scaled = neg_out + (pos_out - neg_out) * cond_scale
return scaled
self.profile_start(model_name, color='blue')
latent = latent.half().cuda()
noise = get_noise(latent).cuda()
sigmas = get_sigmas(self.torch_models[model_name].model_sampling, self.denoising_steps).cuda()
sigmas = sigmas[int(self.denoising_steps * (1 - self.denoising_percentage)):]
conditioning = fix_cond(conditioning)
neg_conditioning = fix_cond(neg_conditioning)
noise_scaled = self.torch_models[model_name].model_sampling.noise_scaling(sigmas[0], noise, latent, max_denoise(sigmas))
extra_args = { "cond": conditioning, "uncond": neg_conditioning, "cond_scale": self.cfg_scale }
latent = sample_euler(cfg_denoiser, noise_scaled, sigmas, extra_args=extra_args)
latent = SD3LatentFormat().process_out(latent)
self.profile_stop(model_name)
return latent
def encode_image(self, model_name='vae_encoder'):
self.input_image = self.input_image.to(self.device)
self.profile_start(model_name, color='orange')
if self.torch_inference or self.torch_fallback[model_name]:
with torch.autocast("cuda", dtype=torch.float16):
latent = self.torch_models[model_name](self.input_image)
else:
latent = self.runEngine(model_name, {'images': self.input_image})['latent']
latent = SD3LatentFormat().process_in(latent)
self.profile_stop(model_name)
return latent
def decode_latent(self, latent, model_name='vae_decoder'):
self.profile_start(model_name, color='red')
if self.torch_inference or self.torch_fallback[model_name]:
with torch.autocast("cuda", dtype=torch.float16):
image = self.torch_models[model_name](latent)
else:
image = self.runEngine(model_name, {'latent': latent})['images']
image = image.float()
self.profile_stop(model_name)
return image
def infer(
self,
prompt,
negative_prompt,
image_height,
image_width,
warmup=False,
save_image=True,
):
"""
Run the diffusion pipeline.
Args:
prompt (str):
The text prompt to guide image generation.
negative_prompt (str):
The prompt not to guide the image generation.
image_height (int):
Height (in pixels) of the image to be generated. Must be a multiple of 8.
image_width (int):
Width (in pixels) of the image to be generated. Must be a multiple of 8.
warmup (bool):
Indicate if this is a warmup run.
save_image (bool):
Save the generated image (if applicable)
"""
assert len(prompt) == len(negative_prompt)
batch_size = len(prompt)
# Spatial dimensions of latent tensor
latent_height = image_height // 8
latent_width = image_width // 8
if self.generator and self.seed:
self.generator.manual_seed(self.seed)
with torch.inference_mode(), trt.Runtime(TRT_LOGGER):
torch.cuda.synchronize()
e2e_tic = time.perf_counter()
# Initialize Latents
latent = self.initialize_latents(batch_size=batch_size,
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()
# Get Conditionings
conditioning, neg_conditioning = self.encode_prompt(prompt, negative_prompt)
# Denoise
latent = self.denoise_latent(latent, conditioning, neg_conditioning)
# Decode Latents
images = self.decode_latent(latent)
torch.cuda.synchronize()
e2e_toc = time.perf_counter()
walltime_ms = (e2e_toc - e2e_tic) * 1000.
if not warmup:
num_inference_steps = int(self.denoising_steps * self.denoising_percentage)
self.print_summary(num_inference_steps, walltime_ms, batch_size)
if save_image:
# post-process images
images = ((images + 1) * 255 / 2).clamp(0, 255).detach().permute(0, 2, 3, 1).round().type(torch.uint8).cpu().numpy()
self.save_image(images, self.pipeline_type.name.lower(), prompt, self.seed)
return images, walltime_ms
def run(self, prompt, negative_prompt, height, width, batch_size, batch_count, num_warmup_runs, use_cuda_graph, **kwargs):
# Process prompt
if not isinstance(prompt, list):
raise ValueError(f"`prompt` must be of type `str` list, but is {type(prompt)}")
prompt = prompt * batch_size
if not isinstance(negative_prompt, list):
raise ValueError(f"`--negative-prompt` must be of type `str` list, but is {type(negative_prompt)}")
if len(negative_prompt) == 1:
negative_prompt = negative_prompt * batch_size
num_warmup_runs = max(1, num_warmup_runs) if use_cuda_graph else num_warmup_runs
if num_warmup_runs > 0:
print("[I] Warming up ..")
for _ in range(num_warmup_runs):
self.infer(prompt, negative_prompt, height, width, warmup=True, **kwargs)
for _ in range(batch_count):
print("[I] Running StableDiffusion3 pipeline")
if self.nvtx_profile:
cudart.cudaProfilerStart()
self.infer(prompt, negative_prompt, height, width, warmup=False, **kwargs)
if self.nvtx_profile:
cudart.cudaProfilerStop()
demo/Diffusion/demo_diffusion/pipeline/stable_diffusion_pipeline.py
+1101
View File
File diff suppressed because it is too large Load Diff
demo/Diffusion/demo_diffusion/pipeline/stable_video_diffusion_pipeline.py
+751
View File
@@ -0,0 +1,751 @@
#
# Copyright 2024 The HuggingFace Inc. team.
# 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 pathlib
import random
import sys
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.bindings import runtime as cudart
from diffusers.image_processor import VaeImageProcessor
from diffusers.utils.torch_utils import randn_tensor
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,
)
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,
)
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):
"""
Application showcasing the acceleration of Stable Video Diffusion pipelines using NVidia TensorRT.
"""
def __init__(
self,
version='svd-xt-1.1',
pipeline_type=PIPELINE_TYPE.IMG2VID,
min_guidance_scale: float = 1.0,
max_guidance_scale: float = 3.0,
decode_chunk_size: Optional[int] = None,
**kwargs
):
"""
Initializes the Diffusion pipeline.
Args:
version (str):
The version of the pipeline. Should be one of [svd-xt-1.1]
pipeline_type (PIPELINE_TYPE):
Type of current pipeline.
min_guidance_scale (`float`, *optional*, defaults to 1.0):
The minimum guidance scale. Used for the classifier free guidance with first frame.
max_guidance_scale (`float`, *optional*, defaults to 3.0):
The maximum guidance scale. Used for the classifier free guidance with last frame.
`max_guidance_scale = 1` corresponds to doing no classifier free guidance.
decode_chunk_size (`int`, *optional*):
The number of frames to decode at a time. The higher the chunk size, the higher the temporal consistency
between frames, but also the higher the memory consumption. By default, the decoder will decode all frames at once
for maximal quality. Reduce `decode_chunk_size` to reduce memory usage.
"""
super().__init__(
version=version,
pipeline_type=pipeline_type,
**kwargs
)
self.min_guidance_scale = min_guidance_scale
self.max_guidance_scale = max_guidance_scale
self.do_classifier_free_guidance = max_guidance_scale > 1
# FIXME vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1)
self.vae_scale_factor = 8
# FIXME num_frames = self.config.num_frames
select_num_frames = {
'svd-xt-1.1': 25,
}
self.num_frames = select_num_frames.get(version, 14)
# TODO decode_chunk_size from args
self.decode_chunk_size = 8 if not decode_chunk_size else decode_chunk_size
# TODO: scaling_factor = vae.config.scaling_factor
self.scaling_factor = 0.18215
# TODO user configurable cuda_device_id
cuda_device_id = 0
properties = cudart.cudaGetDeviceProperties(cuda_device_id)
if properties[0] != 0:
total_device_count = cudart.cudaGetDeviceCount()[1]
raise ValueError(f"Failed to get device properties for device {cuda_device_id}, total device count: {total_device_count}")
vram_size = properties[1].totalGlobalMem
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.")
if self.use_cuda_graph and self.low_vram:
print("[W] WARNING CUDA graph disabled in low VRAM mode.")
self.use_cuda_graph = False
self.config = {}
if self.pipeline_type.is_img2vid():
self.config['clip_vis_torch_fallback'] = True
self.config['clip_imgfe_torch_fallback'] = True
self.config['vae_temp_torch_fallback'] = True
# initialized in loadEngines()
self.max_shared_device_memory_size = 0
def loadResources(self, image_height, image_width, batch_size, seed):
# Initialize noise generator
self.seed = seed
self.generator = torch.Generator(device="cuda").manual_seed(seed) if seed else None
# Create CUDA events and stream
for stage in ['clip', 'denoise', 'vae', 'vae_encoder']:
self.events[stage] = [cudart.cudaEventCreate()[1], cudart.cudaEventCreate()[1]]
self.stream = cudart.cudaStreamCreate()[1]
# Allocate shared device memory for TensorRT engines
if not self.low_vram and not self.torch_inference:
for model_name in self.models.keys():
if not self.torch_fallback[model_name]:
self.max_shared_device_memory_size = max(self.max_shared_device_memory_size, self.engine[model_name].engine.device_memory_size_v2)
self.shared_device_memory = cudart.cudaMalloc(self.max_shared_device_memory_size)[1]
# Activate TensorRT engines
for model_name in self.models.keys():
if not self.torch_fallback[model_name]:
self.engine[model_name].activate(device_memory=self.shared_device_memory)
alloc_shape = self.models[model_name].get_shape_dict(batch_size, image_height, image_width)
self.engine[model_name].allocate_buffers(shape_dict=alloc_shape, device=self.device)
def loadEngines(
self,
engine_dir,
framework_model_dir,
onnx_dir,
onnx_opset,
opt_batch_size,
opt_image_height,
opt_image_width,
static_batch=False,
static_shape=True,
enable_refit=False,
enable_all_tactics=False,
timing_cache=None,
fp8=False,
quantization_level=0.0,
calibration_size=32,
calib_batch_size=2,
**_kwargs,
):
"""
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):
Batch size to optimize for during engine building.
opt_image_height (int):
Image height to optimize for during engine building. Must be a multiple of 8.
opt_image_width (int):
Image width to optimize for during engine building. Must be a multiple of 8.
static_batch (bool):
Build engine only for specified opt_batch_size.
static_shape (bool):
Build engine only for specified opt_image_height & opt_image_width. Default = True.
enable_refit (bool):
Build engines with refit option enabled.
enable_all_tactics (bool):
Enable all tactic sources during TensorRT engine builds.
timing_cache (str):
Path to the timing cache to speed up TensorRT build.
fp8 (bool):
Whether to quantize to fp8 format or not.
quantization_level (float):
Controls which layers to quantize.
calibration_size (int):
The number of steps to use for calibrating the model for quantization.
calib_batch_size (int):
The batch size to use for calibration. Defaults to 2.
"""
# 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)
# Load pipeline models
models_args = {'version': self.version, 'pipeline': self.pipeline_type, 'device': self.device,
'hf_token': self.hf_token, 'verbose': self.verbose, 'framework_model_dir': framework_model_dir,
'max_batch_size': self.max_batch_size}
if 'clip-vis' in self.stages:
self.models['clip-vis'] = CLIPVisionWithProjModel(**models_args, subfolder='image_encoder')
if 'clip-imgfe' in self.stages:
self.models['clip-imgfe'] = CLIPImageProcessorModel(**models_args, subfolder='feature_extractor')
if 'unet-temp' in self.stages:
self.models['unet-temp'] = UNetTemporalModel(**models_args, fp16=True, fp8=fp8, num_frames=self.num_frames, do_classifier_free_guidance=self.do_classifier_free_guidance)
if 'vae-temp' in self.stages:
self.models['vae-temp'] = VAEDecTemporalModel(**models_args, decode_chunk_size=self.decode_chunk_size)
self.image_processor = VaeImageProcessor(vae_scale_factor=self.vae_scale_factor)
# Configure pipeline models to load
model_names = self.models.keys()
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]))
onnx_path = dict(zip(model_names, [self.getOnnxPath(model_name, onnx_dir, opt=False) for model_name in model_names]))
onnx_opt_path = dict(zip(model_names, [self.getOnnxPath(model_name, onnx_dir) for model_name in model_names]))
engine_path = dict(zip(model_names, [self.getEnginePath(model_name, engine_dir) for model_name in model_names]))
do_engine_refit = dict(zip(model_names, [enable_refit and model_name.startswith('unet') for model_name in model_names]))
# Quantization.
model_suffix = dict(zip(model_names, ['' for model_name in model_names]))
use_fp8 = dict.fromkeys(model_names, False)
if fp8:
model_name = "unet-temp"
use_fp8[model_name] = True
model_suffix[model_name] += f"-fp8.l{quantization_level}.bs2.s{self.denoising_steps}.c{calibration_size}"
onnx_path = { model_name : self.getOnnxPath(model_name, onnx_dir, opt=False, suffix=model_suffix[model_name]) for model_name in model_names }
onnx_opt_path = { model_name : self.getOnnxPath(model_name, onnx_dir, suffix=model_suffix[model_name]) for model_name in model_names }
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]:
continue
do_export_onnx = not os.path.exists(engine_path[model_name]) and not os.path.exists(onnx_opt_path[model_name])
do_export_weights_map = weights_map_path[model_name] and not os.path.exists(weights_map_path[model_name])
if do_export_onnx or do_export_weights_map:
if use_fp8[model_name]:
pipeline = obj.get_pipeline()
model = pipeline.unet
state_dict_path = self.getStateDictPath(model_name, onnx_dir, suffix=model_suffix[model_name])
if not os.path.exists(state_dict_path):
# Load calibration images
print(f"[I] Calibrated weights not found, generating {state_dict_path}")
root_dir = os.path.dirname(os.path.abspath(sys.modules["__main__"].__file__))
calibration_image_folder = os.path.join(root_dir, "calibration_data", "calibration-images")
calibration_image_list = load_calibration_images(calibration_image_folder)
print("Number of images loaded:", len(calibration_image_list))
# TODO check size > calibration_size
def do_calibrate(pipeline, calibration_images, **kwargs):
for i_th, image in enumerate(calibration_images):
if i_th >= kwargs["calib_size"]:
return
pipeline(
image=image,
num_inference_steps=kwargs["n_steps"],
).frames[0]
def forward_loop(model):
pipeline.unet = model
do_calibrate(
pipeline=pipeline,
calibration_images=calibration_image_list,
calib_size=calibration_size // calib_batch_size,
n_steps=self.denoising_steps,
)
print(f"[I] Performing calibration for {calibration_size} steps.")
if use_fp8[model_name]:
quant_config = SD_FP8_FP16_DEFAULT_CONFIG
check_lora(model)
mtq.quantize(model, quant_config, forward_loop)
mto.save(model, state_dict_path)
else:
mto.restore(model, state_dict_path)
print(f"[I] Generating quantized ONNX model: {onnx_opt_path[model_name]}")
if not os.path.exists(onnx_path[model_name]):
"""
Error: Torch bug, ONNX export failed due to unknown kernel shape in QuantConv3d.
TRT_FP8QuantizeLinear and TRT_FP8DequantizeLinear operations in UNetSpatioTemporalConditionModel for svd
cause issues. Inputs on different devices (CUDA vs CPU) may contribute to the problem.
"""
quantize_lvl(self.version, model, quantization_level, enable_conv_3d=False)
mtq.disable_quantizer(model, filter_func)
if use_fp8[model_name]:
generate_fp8_scales(model)
else:
model = None
obj.export_onnx(onnx_path[model_name], onnx_opt_path[model_name], onnx_opset, opt_image_height, opt_image_width, custom_model=model, static_shape=static_shape)
else:
obj.export_onnx(onnx_path[model_name], onnx_opt_path[model_name], onnx_opset, opt_image_height, opt_image_width)
# Clean model cache
torch.cuda.empty_cache()
# Build TensorRT engines
for model_name, obj in self.models.items():
if self.torch_fallback[model_name]:
continue
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],
input_profile=obj.get_input_profile(
opt_batch_size, opt_image_height, opt_image_width,
static_batch=static_batch, static_shape=static_shape
),
enable_refit=do_engine_refit[model_name],
enable_all_tactics=enable_all_tactics,
timing_cache=timing_cache,
update_output_names=update_output_names,
native_instancenorm=False)
self.engine[model_name] = engine
# Load TensorRT engines
for model_name in self.models.keys():
if not self.torch_fallback[model_name]:
self.engine[model_name].load()
def activateEngines(self, model_name, alloc_shape=None):
if not self.torch_fallback[model_name]:
device_memory_update = self.low_vram and not self.shared_device_memory
if device_memory_update:
assert not self.use_cuda_graph
# Reclaim GPU memory from torch cache
torch.cuda.empty_cache()
self.shared_device_memory = cudart.cudaMalloc(self.max_shared_device_memory_size)[1]
# Create TensorRT execution context
if not self.engine[model_name].context:
assert not self.use_cuda_graph
self.engine[model_name].activate(device_memory=self.shared_device_memory)
if device_memory_update:
self.engine[model_name].reactivate(device_memory=self.shared_device_memory)
if alloc_shape and not self.engine[model_name].tensors:
assert not self.use_cuda_graph
self.engine[model_name].allocate_buffers(shape_dict=alloc_shape, device=self.device)
else:
# Load torch model
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):
if not release_model:
return
if not self.torch_fallback[model_name]:
assert not self.use_cuda_graph
self.engine[model_name].deallocate_buffers()
self.engine[model_name].deactivate()
# Shared device memory deallocated only in low VRAM mode
if self.low_vram and self.shared_device_memory:
cudart.cudaFree(self.shared_device_memory)
self.shared_device_memory = None
else:
del self.torch_models[model_name]
def print_summary(self, denoising_steps, walltime_ms, batch_size, num_frames):
print('|-----------------|--------------|')
print('| {:^15} | {:^12} |'.format('Module', 'Latency'))
print('|-----------------|--------------|')
print('| {:^15} | {:>9.2f} ms |'.format('VAE-Enc', cudart.cudaEventElapsedTime(self.events['vae_encoder'][0], self.events['vae_encoder'][1])[1]))
print('| {:^15} | {:>9.2f} ms |'.format('CLIP', cudart.cudaEventElapsedTime(self.events['clip'][0], self.events['clip'][1])[1]))
print('| {:^15} | {:>9.2f} ms |'.format('UNet'+('+CNet' if self.pipeline_type.is_controlnet() else '')+' x '+str(denoising_steps), cudart.cudaEventElapsedTime(self.events['denoise'][0], self.events['denoise'][1])[1]))
print('| {:^15} | {:>9.2f} ms |'.format('VAE-Dec', cudart.cudaEventElapsedTime(self.events['vae'][0], self.events['vae'][1])[1]))
print('|-----------------|--------------|')
print('| {:^15} | {:>9.2f} ms |'.format('Pipeline', walltime_ms))
print('|-----------------|--------------|')
print('Throughput: {:.5f} videos/min ({} frames)'.format(batch_size*60000./walltime_ms, num_frames))
def save_video(self, frames, pipeline, seed):
video_name_prefix = '-'.join([pipeline, 'fp16', str(seed), str(random.randint(1000,9999))])
video_name_suffix = 'torch' if self.torch_inference else 'trt'
video_path = video_name_prefix+'-'+video_name_suffix+'.gif'
print(f"Saving video to: {video_path}")
frames[0].save(os.path.join(self.output_dir, video_path), save_all=True,optimize=False, append_images=frames[1:], loop=0)
def _encode_image(self, image, num_videos_per_prompt, do_classifier_free_guidance):
dtype = next(self.torch_models['clip-vis'].parameters()).dtype
if not isinstance(image, torch.Tensor):
image = self.image_processor.pil_to_numpy(image)
image = self.image_processor.numpy_to_pt(image)
# 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 = image_module.resize_with_antialiasing(image, (224, 224))
image = (image + 1.0) / 2.0
# Normalize the image with for CLIP input
image = self.torch_models['clip-imgfe'](
images=image,
do_normalize=True,
do_center_crop=False,
do_resize=False,
do_rescale=False,
return_tensors="pt",
).pixel_values
image = image.to(device=self.device, dtype=dtype)
image_embeddings = self.torch_models['clip-vis'](image).image_embeds
image_embeddings = image_embeddings.unsqueeze(1)
# duplicate image embeddings for each generation per prompt, using mps friendly method
bs_embed, seq_len, _ = image_embeddings.shape
image_embeddings = image_embeddings.repeat(1, num_videos_per_prompt, 1)
image_embeddings = image_embeddings.view(bs_embed * num_videos_per_prompt, seq_len, -1)
if do_classifier_free_guidance:
negative_image_embeddings = torch.zeros_like(image_embeddings)
# For classifier free guidance, we need to do two forward passes.
# Here we concatenate the unconditional and text embeddings into a single batch
# to avoid doing two forward passes
image_embeddings = torch.cat([negative_image_embeddings, image_embeddings])
return image_embeddings
def _encode_vae_image(
self,
image: torch.Tensor,
device,
num_videos_per_prompt,
do_classifier_free_guidance,
):
image = image.to(device=device)
image_latents = self.torch_models['vae-temp'].encode(image).latent_dist.mode()
if do_classifier_free_guidance:
negative_image_latents = torch.zeros_like(image_latents)
# For classifier free guidance, we need to do two forward passes.
# Here we concatenate the unconditional and text embeddings into a single batch
# to avoid doing two forward passes
image_latents = torch.cat([negative_image_latents, image_latents])
# duplicate image_latents for each generation per prompt, using mps friendly method
image_latents = image_latents.repeat(num_videos_per_prompt, 1, 1, 1)
return image_latents
def _get_add_time_ids(
self,
fps,
motion_bucket_id,
noise_aug_strength,
dtype,
batch_size,
num_videos_per_prompt,
do_classifier_free_guidance,
):
add_time_ids = [fps, motion_bucket_id, noise_aug_strength]
add_time_ids = torch.tensor([add_time_ids], dtype=dtype)
add_time_ids = add_time_ids.repeat(batch_size * num_videos_per_prompt, 1)
if do_classifier_free_guidance:
add_time_ids = torch.cat([add_time_ids, add_time_ids])
return add_time_ids
def prepare_latents(
self,
batch_size,
num_frames,
num_channels_latents,
height,
width,
dtype,
device,
latents=None,
):
shape = (
batch_size,
num_frames,
num_channels_latents // 2,
height // self.vae_scale_factor,
width // self.vae_scale_factor,
)
if latents is None:
latents = randn_tensor(shape, generator=self.generator, device=device, dtype=dtype)
else:
latents = latents.to(device)
# scale the initial noise by the standard deviation required by the scheduler
latents = latents * self.scheduler.init_noise_sigma
return latents
def decode_latents(self, latents, num_frames, decode_chunk_size):
# [batch, frames, channels, height, width] -> [batch*frames, channels, height, width]
latents = latents.flatten(0, 1)
latents = 1 / self.scaling_factor * latents
# decode decode_chunk_size frames at a time to avoid OOM
frames = []
for i in range(0, latents.shape[0], decode_chunk_size):
num_frames_in = latents[i : i + decode_chunk_size].shape[0]
# TODO only pass num_frames_in if it's expected
if self.torch_fallback['vae-temp']:
frame = self.torch_models['vae-temp'].decode(latents[i : i + decode_chunk_size], num_frames=num_frames_in).sample
else:
params = {
'latent': latents[i : i + decode_chunk_size],
# FIXME segfault
#'num_frames_in': torch.Tensor([num_frames_in]).to(device=latents.device, dtype=torch.int64),
}
frame = self.runEngine('vae-temp', params)['frames']
frames.append(frame)
frames = torch.cat(frames, dim=0)
# [batch*frames, channels, height, width] -> [batch, channels, frames, height, width]
frames = frames.reshape(-1, num_frames, *frames.shape[1:]).permute(0, 2, 1, 3, 4)
# we always cast to float32 as this does not cause significant overhead and is compatible with bfloat16
frames = frames.float()
return frames
def infer(
self,
input_image,
image_height: int,
image_width: int,
fps: int = 7,
motion_bucket_id: int = 127,
noise_aug_strength: int = 0.02,
num_videos_per_prompt: Optional[int] = 1,
warmup: bool = False,
save_video: bool = True,
):
"""
Run the video diffusion pipeline.
Args:
input_image (image):
Input image used to initialize the latents.
image_height (int):
Height (in pixels) of the image to be generated. Must be a multiple of 8.
image_width (int):
Width (in pixels) of the image to be generated. Must be a multiple of 8.
fps (`int`, *optional*, defaults to 7):
Frames per second. The rate at which the generated images shall be exported to a video after generation.
Note that Stable Diffusion Video's UNet was micro-conditioned on fps-1 during training.
motion_bucket_id (`int`, *optional*, defaults to 127):
The motion bucket ID. Used as conditioning for the generation. The higher the number the more motion will be in the video.
noise_aug_strength (`int`, *optional*, defaults to 0.02):
The amount of noise added to the init image, the higher it is the less the video will look like the init image. Increase it for more motion.
num_videos_per_prompt (`int`, *optional*, defaults to 1):
The number of images to generate per prompt.
warmup (bool):
Indicate if this is a warmup run.
save_video (bool):
Save the video image.
"""
if self.generator and self.seed:
self.generator.manual_seed(self.seed)
# TODO
batch_size = 1
# Fast warmup
denoising_steps = 1 if warmup else self.denoising_steps
torch.cuda.synchronize()
e2e_tic = time.perf_counter()
class LoadModelContext:
def __init__(ctx, model_name, alloc_shape=None, release_model=False):
ctx.model_name = model_name
ctx.release_model = release_model
ctx.alloc_shape = alloc_shape
def __enter__(ctx):
self.activateEngines(ctx.model_name, alloc_shape=ctx.alloc_shape)
def __exit__(ctx, exc_type, exc_val, exc_tb):
self.deactivateEngines(ctx.model_name, release_model=ctx.release_model)
# Release model opportunistically in TensorRT pipeline only in low VRAM mode
release_model = self.low_vram and not self.torch_inference
with torch.inference_mode(), torch.autocast("cuda"), trt.Runtime(TRT_LOGGER):
with LoadModelContext('clip-imgfe', release_model=release_model), LoadModelContext('clip-vis', release_model=release_model):
self.profile_start('clip', color='green')
image_embeddings = self._encode_image(input_image, num_videos_per_prompt, self.do_classifier_free_guidance)
self.profile_stop('clip')
# NOTE Stable Diffusion Video was conditioned on fps - 1
fps = fps - 1
self.profile_start('preprocess', color='pink')
input_image = self.image_processor.preprocess(input_image, height=image_height, width=image_width).to(self.device)
noise = randn_tensor(input_image.shape, generator=self.generator, device=input_image.device, dtype=input_image.dtype)
input_image = input_image + noise_aug_strength * noise
self.profile_stop('preprocess')
# TODO
# assert self.torch_models['vae-temp'].dtype == torch.float32
with LoadModelContext('vae-temp'):
self.profile_start('vae_encoder', color='red')
image_latents = self._encode_vae_image(input_image, self.device, num_videos_per_prompt, self.do_classifier_free_guidance)
image_latents = image_latents.to(image_embeddings.dtype)
self.profile_stop('vae_encoder')
# Repeat the image latents for each frame so we can concatenate them with the noise
# image_latents [batch, channels, height, width] ->[batch, num_frames, channels, height, width]
image_latents = image_latents.unsqueeze(1).repeat(1, self.num_frames, 1, 1, 1)
# Get Added Time IDs
added_time_ids = self._get_add_time_ids(
fps,
motion_bucket_id,
noise_aug_strength,
image_embeddings.dtype,
batch_size,
num_videos_per_prompt,
self.do_classifier_free_guidance,
)
added_time_ids = added_time_ids.to(self.device)
# Prepare timesteps
self.scheduler.set_timesteps(denoising_steps, device=self.device)
timesteps = self.scheduler.timesteps
# Prepare latent variables
latents = self.prepare_latents(
batch_size * num_videos_per_prompt,
self.num_frames,
8, # TODO: num_channels_latents = unet.config.in_channels
image_height,
image_width,
image_embeddings.dtype,
input_image.device,
None, # pre-generated latents
)
# Prepare guidance scale
guidance_scale = torch.linspace(self.min_guidance_scale, self.max_guidance_scale, self.num_frames).unsqueeze(0)
guidance_scale = guidance_scale.to(self.device, latents.dtype)
guidance_scale = guidance_scale.repeat(batch_size * num_videos_per_prompt, 1)
guidance_scale = _append_dims(guidance_scale, latents.ndim)
# Denoising loop
num_warmup_steps = len(timesteps) - denoising_steps * self.scheduler.order
unet_shape_dict = self.models['unet-temp'].get_shape_dict(batch_size, image_height, image_width)
with LoadModelContext('unet-temp', alloc_shape=unet_shape_dict, release_model=release_model), tqdm(total=denoising_steps) as progress_bar:
self.profile_start('denoise', color='blue')
for i, t in enumerate(timesteps):
# expand the latents if we are doing classifier free guidance
latent_model_input = torch.cat([latents] * 2) if self.do_classifier_free_guidance else latents
latent_model_input = self.scheduler.scale_model_input(latent_model_input, t)
# Concatenate image_latents over channels dimention
latent_model_input = torch.cat([latent_model_input, image_latents], dim=2)
# predict the noise residual
if self.torch_fallback['unet-temp']:
noise_pred = self.torch_models['unet-temp'](
latent_model_input,
t,
encoder_hidden_states=image_embeddings,
added_time_ids=added_time_ids,
return_dict=False,
)[0]
else:
params = {
"sample": latent_model_input,
"timestep": t,
"encoder_hidden_states": image_embeddings,
"added_time_ids": added_time_ids,
}
noise_pred = self.runEngine('unet-temp', params)['latent']
# perform guidance
if self.do_classifier_free_guidance:
noise_pred_uncond, noise_pred_cond = noise_pred.chunk(2)
noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_cond - noise_pred_uncond)
# compute the previous noisy sample x_t -> x_t-1
latents = self.scheduler.step(noise_pred, t, latents).prev_sample
if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0):
progress_bar.update()
self.profile_stop('denoise')
with torch.inference_mode(), trt.Runtime(TRT_LOGGER), LoadModelContext('vae-temp'):
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 = image_module.tensor2vid(frames, self.image_processor, output_type="pil")
self.profile_stop('vae')
torch.cuda.synchronize()
if warmup:
return
e2e_toc = time.perf_counter()
walltime_ms = (e2e_toc - e2e_tic) * 1000.
self.print_summary(denoising_steps, walltime_ms, batch_size, len(frames[0]))
if save_video:
self.save_video(frames[0], self.pipeline_type.name.lower(), self.seed)
return frames, walltime_ms
def run(self, input_image, height, width, batch_size, batch_count, num_warmup_runs, use_cuda_graph, **kwargs):
num_warmup_runs = max(1, num_warmup_runs) if use_cuda_graph else num_warmup_runs
if num_warmup_runs > 0:
print("[I] Warming up ..")
for _ in range(num_warmup_runs):
self.infer(input_image, height, width, warmup=True)
for _ in range(batch_count):
print("[I] Running StableDiffusion pipeline")
if self.nvtx_profile:
cudart.cudaProfilerStart()
self.infer(input_image, height, width, warmup=False)
if self.nvtx_profile:
cudart.cudaProfilerStop()
demo/Diffusion/demo_diffusion/pipeline/type.py
+68
View File
@@ -0,0 +1,68 @@
#
# 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()
TXT2VID = 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()
VIDEO2WORLD = 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_txt2vid(self):
return self == self.TXT2VID
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()
def is_video2world(self):
return self == self.VIDEO2WORLD
demo/Diffusion/demo_diffusion/pipeline/wan_pipeline.py
+841
View File
@@ -0,0 +1,841 @@
#
# 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 argparse
import gc
import html
import os
import pathlib
import random
import re
import time
from typing import Any, Callable, Dict, List, Optional, Union
import tensorrt as trt
import torch
from cuda.bindings import runtime as cudart
from diffusers.utils.torch_utils import randn_tensor
from PIL import Image
from tqdm.auto import tqdm
try:
import ftfy
FTFY_AVAILABLE = True
except ImportError:
FTFY_AVAILABLE = False
import demo_diffusion.engine as engine_module
import demo_diffusion.image as image_module
import demo_diffusion.path as path_module
from demo_diffusion.model import (
T5Model,
WanTransformerModel,
AutoencoderKLWanModel,
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)
# Copied from https://github.com/huggingface/diffusers/blob/main/src/diffusers/pipelines/wan/pipeline_wan.py#L78
def basic_clean(text):
if FTFY_AVAILABLE:
text = ftfy.fix_text(text)
text = html.unescape(html.unescape(text))
return text.strip()
# Copied from https://github.com/huggingface/diffusers/blob/main/src/diffusers/pipelines/wan/pipeline_wan.py#L84
def whitespace_clean(text):
text = re.sub(r"\s+", " ", text)
text = text.strip()
return text
# Copied from https://github.com/huggingface/diffusers/blob/main/src/diffusers/pipelines/wan/pipeline_wan.py#L90
def prompt_clean(text):
text = whitespace_clean(basic_clean(text))
return text
class WanPipeline(DiffusionPipeline):
"""
Application showcasing the acceleration of Wan 2.2 T2V pipeline using Nvidia TensorRT.
"""
_callback_tensor_inputs = ["latents", "prompt_embeds", "negative_prompt_embeds"]
def __init__(
self,
dd_path,
version='wan2.2-t2v-a14b',
pipeline_type=PIPELINE_TYPE.TXT2VID,
boundary_ratio: float = 0.875,
guidance_scale: float = 4.0,
guidance_scale_2: float = 3.0,
t5_weight_streaming_budget_percentage=None,
transformer_weight_streaming_budget_percentage=None,
**kwargs
):
"""
Initializes the Wan T2V 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 [wan2.2-t2v-a14b]
pipeline_type (PIPELINE_TYPE):
Type of current pipeline (TXT2VID)
boundary_ratio (float, defaults to 0.875):
Ratio of total timesteps to use as the boundary for switching between transformers
in two-stage denoising. Transformer handles high-noise stages (timesteps >= boundary)
and transformer_2 handles low-noise stages (timesteps < boundary).
Wan 2.2 T2V always uses two-stage denoising.
guidance_scale (float):
Guidance scale for high-noise stage (transformer). Wan default: 4.0
guidance_scale_2 (float):
Guidance scale for low-noise stage (transformer_2). Wan default: 3.0
t5_weight_streaming_budget_percentage (`int`, defaults to None):
Weight streaming budget as a percentage of the size of total streamable weights for the T5 model.
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 models.
"""
super().__init__(
dd_path=dd_path,
version=version,
pipeline_type=pipeline_type,
scheduler="UniPC",
bf16=True,
text_encoder_weight_streaming_budget_percentage=t5_weight_streaming_budget_percentage,
denoiser_weight_streaming_budget_percentage=transformer_weight_streaming_budget_percentage,
**kwargs
)
# Validate boundary_ratio (required for Wan 2.2 two-stage denoising)
if boundary_ratio is None or not (0.0 < boundary_ratio < 1.0):
raise ValueError(
f"`boundary_ratio` must be between 0.0 and 1.0, got {boundary_ratio}"
)
self.boundary_ratio = boundary_ratio
self.guidance_scale = guidance_scale
self.guidance_scale_2 = guidance_scale_2
self.vae_scale_factor_temporal = 4
self.vae_scale_factor_spatial = 8
self.opt_image_height = 720
self.opt_image_width = 1280
self.opt_num_frames = 81
self.max_sequence_length = 512
@classmethod
def FromArgs(cls, args: argparse.Namespace, pipeline_type: PIPELINE_TYPE) -> 'WanPipeline':
"""Factory method to construct a WanPipeline object from parsed arguments."""
MAX_BATCH_SIZE = 1 # Wan always uses batch size 1
DEVICE = "cuda"
dd_path = path_module.resolve_path(
cls.get_model_names(pipeline_type), args, pipeline_type, cls._get_pipeline_uid(args.version)
)
return cls(
dd_path=dd_path,
version=args.version,
pipeline_type=pipeline_type,
boundary_ratio=args.boundary_ratio,
denoising_steps=args.denoising_steps,
guidance_scale=args.guidance_scale,
guidance_scale_2=args.guidance_scale_2,
t5_weight_streaming_budget_percentage=args.t5_ws_percentage if hasattr(args, 't5_ws_percentage') else None,
transformer_weight_streaming_budget_percentage=args.transformer_ws_percentage if hasattr(args, 'transformer_ws_percentage') else None,
max_batch_size=MAX_BATCH_SIZE,
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,
low_vram=args.low_vram,
torch_inference=args.torch_inference,
torch_fallback=args.torch_fallback if hasattr(args, 'torch_fallback') else None,
weight_streaming=args.ws if hasattr(args, 'ws') else False,
)
@classmethod
def get_model_names(cls, pipeline_type: PIPELINE_TYPE, controlnet_type: str = None) -> List[str]:
"""Return a list of model names used by this pipeline.
Overrides:
DiffusionPipeline.get_model_names
"""
return ["text_encoder", "transformer", "transformer_2", "vae_decoder"]
def download_onnx_models(self, model_name: str, model_config: dict[str, Any]) -> None:
raise NotImplementedError(
"Pre-exported Wan ONNX models are not available for download. "
"Export ONNX models locally using the provided export script."
)
def _initialize_models(self, framework_model_dir, int8=False, fp8=False, fp4=False):
self.tokenizer = make_tokenizer(
self.version,
self.pipeline_type,
self.hf_token,
framework_model_dir,
subfolder='tokenizer',
tokenizer_type='t5'
)
models_args = {
'version': self.version,
'pipeline': self.pipeline_type,
'device': self.device,
'hf_token': self.hf_token,
'verbose': self.verbose,
'framework_model_dir': framework_model_dir,
'max_batch_size': 1
}
if "text_encoder" in self.stages:
self.models['text_encoder'] = T5Model(
**models_args,
fp16=False,
bf16=True,
text_maxlen=self.max_sequence_length,
weight_streaming=self.weight_streaming,
weight_streaming_budget_percentage=self.text_encoder_weight_streaming_budget_percentage,
use_attention_mask=True,
)
if "transformer" in self.stages:
self.models['transformer'] = WanTransformerModel(
**models_args,
subfolder='transformer',
fp16=False,
bf16=True,
text_maxlen=self.max_sequence_length,
num_frames=self.opt_num_frames,
height=self.opt_image_height,
width=self.opt_image_width,
weight_streaming=self.weight_streaming,
weight_streaming_budget_percentage=self.denoiser_weight_streaming_budget_percentage,
)
if "transformer_2" in self.stages:
self.models['transformer_2'] = WanTransformerModel(
**models_args,
subfolder='transformer_2',
fp16=False,
bf16=True,
text_maxlen=self.max_sequence_length,
num_frames=self.opt_num_frames,
height=self.opt_image_height,
width=self.opt_image_width,
weight_streaming=self.weight_streaming,
weight_streaming_budget_percentage=self.denoiser_weight_streaming_budget_percentage,
)
if "vae_decoder" in self.stages:
self.models['vae_decoder'] = AutoencoderKLWanModel(
**models_args,
)
self.config['vae_decoder_torch_fallback'] = True
def load_resources(self, image_height, image_width, batch_size, seed):
"""Override to create additional 'denoise' event for combined transformer timing."""
super().load_resources(image_height, image_width, batch_size, seed)
# additional event for combined denoising timing (both transformers)
self.events['denoise'] = [cudart.cudaEventCreate()[1], cudart.cudaEventCreate()[1]]
def print_summary(self, denoising_steps, walltime_ms, batch_size, num_frames):
print("|----------------------|--------------|")
print("| {:^20} | {:^12} |".format("Module", "Latency"))
print("|----------------------|--------------|")
# calculate transformer timings from combined denoise event
total_denoise_time = cudart.cudaEventElapsedTime(self.events['denoise'][0], self.events['denoise'][1])[1]
transformer_steps_map = {
'transformer': self.transformer_steps if (self.transformer_steps > 0 and self.transformer_2_steps > 0) else denoising_steps,
'transformer_2': self.transformer_2_steps if self.transformer_2_steps > 0 else 0
}
for stage in self.stages:
if stage in transformer_steps_map and transformer_steps_map[stage] > 0:
steps = transformer_steps_map[stage]
time_ms = total_denoise_time * (steps / denoising_steps)
stage_label = f"{stage} x {steps}"
elif stage in transformer_steps_map:
continue # skip transformer_2 if unused
else:
time_ms = cudart.cudaEventElapsedTime(self.events[stage][0], self.events[stage][1])[1]
stage_label = stage
print("| {:^20} | {:>9.2f} ms |".format(stage_label, time_ms))
print("|----------------------|--------------|")
print("| {:^20} | {:>9.2f} ms |".format("Pipeline", walltime_ms))
print("|----------------------|--------------|")
print("Throughput: {:.2f} videos/min ({} frames)".format(batch_size * 60000.0 / walltime_ms, num_frames))
def save_video(self, frames, pipeline, prompt, seed):
if isinstance(prompt, list):
prompt_prefix = ''.join(set([p.replace(' ','_')[:10] for p in prompt]))
else:
prompt_prefix = prompt.replace(' ','_')[:10]
seed_str = str(seed) if seed is not None else 'random'
precision = 'bf16' if self.bf16 else 'fp16' if self.fp16 else 'fp32'
video_name_prefix = '-'.join([pipeline, prompt_prefix, precision, seed_str, str(random.randint(1000,9999))])
video_name_suffix = 'torch' if self.torch_inference else 'trt'
video_path = video_name_prefix+'-'+video_name_suffix+'.gif'
full_path = os.path.join(self.output_dir, video_path)
print(f"Saving video to: {full_path}")
frames[0].save(full_path, save_all=True, optimize=False, append_images=frames[1:], loop=0)
# Adapted from https://github.com/huggingface/diffusers/blob/main/src/diffusers/pipelines/wan/pipeline_wan.py#L198
def encode_prompt(
self,
prompt: Union[str, List[str]],
negative_prompt: Optional[Union[str, List[str]]] = None,
do_classifier_free_guidance: bool = True,
num_videos_per_prompt: int = 1,
prompt_embeds: Optional[torch.Tensor] = None,
negative_prompt_embeds: Optional[torch.Tensor] = None,
max_sequence_length: int = 226,
device: Optional[torch.device] = None,
dtype: Optional[torch.dtype] = None,
):
r"""
Encodes the prompt into text encoder hidden states.
Implementation modeled from diffusers Wan pipeline, adapted for TensorRT.
Args:
prompt (`str` or `List[str]`, *optional*):
prompt to be encoded
negative_prompt (`str` or `List[str]`, *optional*):
The prompt or prompts not to guide the image generation. If not defined, one has to pass
`negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
less than `1`).
do_classifier_free_guidance (`bool`, *optional*, defaults to `True`):
Whether to use classifier free guidance or not.
num_videos_per_prompt (`int`, *optional*, defaults to 1):
Number of videos that should be generated per prompt. torch device to place the resulting embeddings on
prompt_embeds (`torch.Tensor`, *optional*):
Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not
provided, text embeddings will be generated from `prompt` input argument.
negative_prompt_embeds (`torch.Tensor`, *optional*):
Pre-generated negative text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt
weighting. If not provided, negative_prompt_embeds will be generated from `negative_prompt` input
argument.
max_sequence_length (`int`, defaults to `226`):
Maximum sequence length for text encoder.
device: (`torch.device`, *optional*):
torch device
dtype: (`torch.dtype`, *optional*):
torch dtype
"""
self.profile_start('text_encoder', color='green')
device = device or self._execution_device
prompt = [prompt] if isinstance(prompt, str) else prompt
if prompt is not None:
batch_size = len(prompt)
else:
batch_size = prompt_embeds.shape[0]
if prompt_embeds is None:
prompt_embeds = self._get_t5_prompt_embeds(
prompt=prompt,
num_videos_per_prompt=num_videos_per_prompt,
max_sequence_length=max_sequence_length,
device=device,
dtype=dtype,
)
if do_classifier_free_guidance and negative_prompt_embeds is None:
negative_prompt = negative_prompt or ""
negative_prompt = batch_size * [negative_prompt] if isinstance(negative_prompt, str) else negative_prompt
if prompt is not None and type(prompt) is not type(negative_prompt):
raise TypeError(
f"`negative_prompt` should be the same type to `prompt`, but got {type(negative_prompt)} !="
f" {type(prompt)}."
)
elif batch_size != len(negative_prompt):
raise ValueError(
f"`negative_prompt`: {negative_prompt} has batch size {len(negative_prompt)}, but `prompt`:"
f" {prompt} has batch size {batch_size}. Please make sure that passed `negative_prompt` matches"
" the batch size of `prompt`."
)
negative_prompt_embeds = self._get_t5_prompt_embeds(
prompt=negative_prompt,
num_videos_per_prompt=num_videos_per_prompt,
max_sequence_length=max_sequence_length,
device=device,
dtype=dtype,
)
self.profile_stop('text_encoder')
return prompt_embeds, negative_prompt_embeds
def denoise_latents(
self,
latents: torch.Tensor,
prompt_embeds: torch.Tensor,
negative_prompt_embeds: Optional[torch.Tensor],
timesteps: torch.Tensor,
guidance_scale: float,
guidance_scale_2: float,
transformer_dtype: torch.dtype,
num_warmup_steps: int,
attention_kwargs: Optional[Dict[str, Any]] = None,
callback_on_step_end: Optional[Callable] = None,
callback_on_step_end_tensor_inputs: Optional[List[str]] = None,
warmup: bool = False,
) -> torch.Tensor:
boundary_timestep = self.boundary_ratio * self.scheduler.config.num_train_timesteps
self.profile_start('denoise', color='blue')
timestep_stages = []
for i, t in enumerate(timesteps):
if t >= boundary_timestep:
timestep_stages.append((i, t, 'transformer', guidance_scale))
else:
timestep_stages.append((i, t, 'transformer_2', guidance_scale_2))
stage_groups = []
if timestep_stages:
current_group = {
'transformer': timestep_stages[0][2],
'guidance_scale': timestep_stages[0][3],
'timesteps': [(timestep_stages[0][0], timestep_stages[0][1])]
}
for i, t, transformer_name, gs in timestep_stages[1:]:
if transformer_name == current_group['transformer']:
current_group['timesteps'].append((i, t))
else:
stage_groups.append(current_group)
current_group = {
'transformer': transformer_name,
'guidance_scale': gs,
'timesteps': [(i, t)]
}
stage_groups.append(current_group)
self.transformer_steps = sum(len(g['timesteps']) for g in stage_groups if g['transformer'] == 'transformer')
self.transformer_2_steps = sum(len(g['timesteps']) for g in stage_groups if g['transformer'] == 'transformer_2')
with tqdm(total=len(timesteps)) as progress_bar:
for stage_group in stage_groups:
transformer_name = stage_group['transformer']
current_guidance_scale = stage_group['guidance_scale']
with self.model_memory_manager([transformer_name], low_vram=self.low_vram):
for step_index, t in stage_group['timesteps']:
latent_model_input = latents.to(transformer_dtype)
timestep = t.expand(latent_model_input.shape[0])
if self.torch_inference or self.torch_fallback[transformer_name]:
current_model = self.torch_models[transformer_name]
noise_pred_cond = current_model(
hidden_states=latent_model_input,
timestep=timestep,
encoder_hidden_states=prompt_embeds,
attention_kwargs=attention_kwargs,
return_dict=False,
)[0]
if self.do_classifier_free_guidance:
noise_pred_uncond = current_model(
hidden_states=latent_model_input,
timestep=timestep,
encoder_hidden_states=negative_prompt_embeds,
attention_kwargs=attention_kwargs,
return_dict=False,
)[0]
noise_pred = noise_pred_uncond + current_guidance_scale * (noise_pred_cond - noise_pred_uncond)
else:
noise_pred = noise_pred_cond
else:
if self.do_classifier_free_guidance:
params_cond = {
"hidden_states": latent_model_input,
"timestep": timestep,
"encoder_hidden_states": prompt_embeds,
}
# conditional engine call
output_cond = self.run_engine(transformer_name, params_cond)['denoised_latents']
noise_pred_cond = output_cond.clone()
params_uncond = {
"hidden_states": latent_model_input,
"timestep": timestep,
"encoder_hidden_states": negative_prompt_embeds,
}
# unconditional engine call
output_uncond = self.run_engine(transformer_name, params_uncond)['denoised_latents']
noise_pred_uncond = output_uncond.clone()
# Apply classifier-free guidance
noise_pred = noise_pred_uncond + current_guidance_scale * (noise_pred_cond - noise_pred_uncond)
else:
# No CFG
params = {
"hidden_states": latent_model_input,
"timestep": timestep,
"encoder_hidden_states": prompt_embeds,
}
noise_pred = self.run_engine(transformer_name, params)['denoised_latents']
latents = self.scheduler.step(noise_pred, t, latents, return_dict=False)[0]
if callback_on_step_end is not None:
callback_kwargs = {}
for k in callback_on_step_end_tensor_inputs:
callback_kwargs[k] = locals()[k]
callback_outputs = callback_on_step_end(self, step_index, t, callback_kwargs)
latents = callback_outputs.pop("latents", latents)
prompt_embeds = callback_outputs.pop("prompt_embeds", prompt_embeds)
negative_prompt_embeds = callback_outputs.pop("negative_prompt_embeds", negative_prompt_embeds)
if step_index == len(timesteps) - 1 or ((step_index + 1) > num_warmup_steps and (step_index + 1) % self.scheduler.order == 0):
progress_bar.update()
self.profile_stop('denoise')
return latents
# Adapted from https://github.com/huggingface/diffusers/blob/main/src/diffusers/pipelines/wan/pipeline_wan.py#L324
def prepare_latents(
self,
batch_size: int,
num_channels_latents: int = 16,
height: int = 720,
width: int = 1280,
num_frames: int = 81,
dtype: Optional[torch.dtype] = None,
device: Optional[torch.device] = None,
generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
latents: Optional[torch.Tensor] = None,
) -> torch.Tensor:
if latents is not None:
return latents.to(device=device, dtype=dtype)
num_latent_frames = (num_frames - 1) // self.vae_scale_factor_temporal + 1
shape = (
batch_size,
num_channels_latents,
num_latent_frames,
int(height) // self.vae_scale_factor_spatial,
int(width) // self.vae_scale_factor_spatial,
)
if isinstance(generator, list) and len(generator) != batch_size:
raise ValueError(
f"You have passed a list of generators of length {len(generator)}, but requested an effective batch"
f" size of {batch_size}. Make sure the batch size matches the length of the generators."
)
if generator is None and hasattr(self, 'generator'):
generator = self.generator
latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype)
return latents
# Copied from https://github.com/huggingface/diffusers/blob/main/src/diffusers/pipelines/wan/pipeline_wan.py#L279
def check_inputs(
self,
prompt,
negative_prompt,
height,
width,
prompt_embeds=None,
negative_prompt_embeds=None,
callback_on_step_end_tensor_inputs=None,
guidance_scale_2=None,
):
if height % 16 != 0 or width % 16 != 0:
raise ValueError(f"`height` and `width` have to be divisible by 16 but are {height} and {width}.")
if callback_on_step_end_tensor_inputs is not None:
pass
if prompt is not None and prompt_embeds is not None:
raise ValueError(
f"Cannot forward both `prompt`: {prompt} and `prompt_embeds`: {prompt_embeds}. Please make sure to"
" only forward one of the two."
)
elif negative_prompt is not None and negative_prompt_embeds is not None:
raise ValueError(
f"Cannot forward both `negative_prompt`: {negative_prompt} and `negative_prompt_embeds`: {negative_prompt_embeds}. Please make sure to"
" only forward one of the two."
)
elif prompt is None and prompt_embeds is None:
raise ValueError(
"Provide either `prompt` or `prompt_embeds`. Cannot leave both `prompt` and `prompt_embeds` undefined."
)
elif prompt is not None and (not isinstance(prompt, str) and not isinstance(prompt, list)):
raise ValueError(f"`prompt` has to be of type `str` or `list` but is {type(prompt)}")
elif negative_prompt is not None and (
not isinstance(negative_prompt, str) and not isinstance(negative_prompt, list)
):
raise ValueError(f"`negative_prompt` has to be of type `str` or `list` but is {type(negative_prompt)}")
# Adapted from https://github.com/huggingface/diffusers/blob/main/src/diffusers/pipelines/wan/pipeline_wan.py#L157
def _get_t5_prompt_embeds(
self,
prompt: Union[str, List[str]] = None,
num_videos_per_prompt: int = 1,
max_sequence_length: int = 226,
device: Optional[torch.device] = None,
dtype: Optional[torch.dtype] = None,
encoder: str = "text_encoder",
):
device = device or self._execution_device
dtype = dtype or (self.torch_models[encoder].dtype if self.torch_fallback.get(encoder) and encoder in self.torch_models else torch.float32)
prompt = [prompt] if isinstance(prompt, str) else prompt
prompt = [prompt_clean(u) for u in prompt]
batch_size = len(prompt)
text_inputs = self.tokenizer(
prompt,
padding="max_length",
max_length=max_sequence_length,
truncation=True,
add_special_tokens=True,
return_attention_mask=True,
return_tensors="pt",
)
text_input_ids, mask = text_inputs.input_ids, text_inputs.attention_mask
seq_lens = mask.gt(0).sum(dim=1).long()
if self.torch_inference or self.torch_fallback[encoder]:
outputs = self.torch_models[encoder](text_input_ids.to(device), mask.to(device))
prompt_embeds = outputs.last_hidden_state.clone()
else:
outputs = self.run_engine(encoder, {
"input_ids": text_input_ids.to(device),
"attention_mask": mask.to(device)
})
prompt_embeds = outputs['text_embeddings'].clone()
prompt_embeds = prompt_embeds.to(dtype=dtype, device=device)
prompt_embeds = [u[:v] for u, v in zip(prompt_embeds, seq_lens)]
prompt_embeds = torch.stack(
[torch.cat([u, u.new_zeros(max_sequence_length - u.size(0), u.size(1))]) for u in prompt_embeds], dim=0
)
_, seq_len, _ = prompt_embeds.shape
prompt_embeds = prompt_embeds.repeat(1, num_videos_per_prompt, 1)
prompt_embeds = prompt_embeds.view(batch_size * num_videos_per_prompt, seq_len, -1)
return prompt_embeds
@property
def _execution_device(self):
return self.device
@property
def do_classifier_free_guidance(self):
return self.guidance_scale > 1.0
@property
def num_timesteps(self):
return self._num_timesteps
def decode_latents(self, latents, num_frames):
self.profile_start('vae_decoder', color='red')
vae_config = self.models['vae_decoder'].config
z_dim = vae_config.get("z_dim", 16)
vae_dtype = torch.float32
latents = latents.to(vae_dtype)
latents_mean = (
torch.tensor(vae_config.get("latents_mean"))
.view(1, z_dim, 1, 1, 1)
.to(latents.device, latents.dtype)
)
latents_std = 1.0 / torch.tensor(vae_config.get("latents_std")).view(1, z_dim, 1, 1, 1).to(
latents.device, latents.dtype
)
latents = latents / latents_std + latents_mean
frames = self.torch_models['vae_decoder'].decode(latents, return_dict=False)[0]
self.profile_stop('vae_decoder')
return frames
def postprocess(self, video: torch.Tensor, output_type: str = "pil"):
# Convert [F, C, H, W] -> [F, H, W, C]
video = video.permute(0, 2, 3, 1)
# Convert to list of PIL Images
video = (video + 1.0) / 2.0
video = torch.clamp(video, 0.0, 1.0)
video = (video * 255.0).to(torch.uint8).cpu().numpy()
pil_frames = [Image.fromarray(frame) for frame in video]
return pil_frames
def infer(
self,
prompt: Union[str, List[str]],
negative_prompt: Optional[Union[str, List[str]]] = None,
height: int = 720,
width: int = 1280,
num_frames: int = 81,
num_inference_steps: int = 40,
num_videos_per_prompt: int = 1,
generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
latents: Optional[torch.Tensor] = None,
prompt_embeds: Optional[torch.Tensor] = None,
negative_prompt_embeds: Optional[torch.Tensor] = None,
output_type: str = "pil",
attention_kwargs: Optional[Dict[str, Any]] = None,
callback_on_step_end: Optional[Union[Callable, Any]] = None,
callback_on_step_end_tensor_inputs: List[str] = ["latents"],
max_sequence_length: int = 512,
warmup: bool = False,
save_video: bool = True,
):
"""
Run the Wan text-to-video diffusion pipeline.
"""
self.check_inputs(
prompt,
negative_prompt,
height,
width,
prompt_embeds,
negative_prompt_embeds,
callback_on_step_end_tensor_inputs,
self.guidance_scale_2,
)
if num_frames % self.vae_scale_factor_temporal != 1:
print(f"[W] `num_frames - 1` has to be divisible by {self.vae_scale_factor_temporal}. Rounding to the nearest number.")
num_frames = num_frames // self.vae_scale_factor_temporal * self.vae_scale_factor_temporal + 1
num_frames = max(num_frames, 1)
device = self._execution_device
batch_size = 1
with torch.inference_mode(), trt.Runtime(TRT_LOGGER):
torch.cuda.synchronize()
e2e_tic = time.perf_counter()
with self.model_memory_manager(["text_encoder"], low_vram=self.low_vram):
prompt_embeds, negative_prompt_embeds = self.encode_prompt(
prompt=prompt,
negative_prompt=negative_prompt,
do_classifier_free_guidance=self.do_classifier_free_guidance,
num_videos_per_prompt=num_videos_per_prompt,
prompt_embeds=prompt_embeds,
negative_prompt_embeds=negative_prompt_embeds,
max_sequence_length=max_sequence_length,
device=device,
)
transformer_dtype = torch.bfloat16
prompt_embeds = prompt_embeds.to(transformer_dtype)
if negative_prompt_embeds is not None:
negative_prompt_embeds = negative_prompt_embeds.to(transformer_dtype)
self.scheduler.set_timesteps(num_inference_steps, device=self.device)
timesteps = self.scheduler.timesteps
num_channels_latents = self.models["transformer"].config.get("in_channels", 16)
latents = self.prepare_latents(
batch_size * num_videos_per_prompt,
num_channels_latents,
height,
width,
num_frames,
torch.float32,
device,
generator,
latents,
)
num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order
self._num_timesteps = len(timesteps)
latents = self.denoise_latents(
latents=latents,
prompt_embeds=prompt_embeds,
negative_prompt_embeds=negative_prompt_embeds,
timesteps=timesteps,
guidance_scale=self.guidance_scale,
guidance_scale_2=self.guidance_scale_2,
transformer_dtype=transformer_dtype,
num_warmup_steps=num_warmup_steps,
attention_kwargs=attention_kwargs,
callback_on_step_end=callback_on_step_end,
callback_on_step_end_tensor_inputs=callback_on_step_end_tensor_inputs,
warmup=warmup,
)
with self.model_memory_manager(["vae_decoder"], low_vram=self.low_vram):
video_raw = self.decode_latents(latents, num_frames)
video = image_module.tensor2vid(video_raw, self, output_type="pil")
torch.cuda.synchronize()
e2e_toc = time.perf_counter()
walltime_ms = (e2e_toc - e2e_tic) * 1000.0
if not warmup:
self.print_summary(num_inference_steps, walltime_ms, batch_size, num_frames)
if save_video:
self.save_video(video[0], self.pipeline_type.name.lower(), prompt, self.seed)
return video, walltime_ms
def run(self, prompt, height, width, num_frames, batch_size, batch_count, num_warmup_runs, use_cuda_graph, **kwargs):
if self.low_vram and self.use_cuda_graph:
print("[W] Using low_vram, use_cuda_graph will be disabled")
self.use_cuda_graph = False
num_warmup_runs = max(1, num_warmup_runs) if use_cuda_graph else num_warmup_runs
if num_warmup_runs > 0:
print("[I] Warming up ..")
for _ in range(num_warmup_runs):
self.infer(prompt, height=height, width=width, num_frames=num_frames, warmup=True, **kwargs)
for _ in range(batch_count):
print("[I] Running Wan T2V pipeline")
if self.nvtx_profile:
cudart.cudaProfilerStart()
self.infer(prompt, height=height, width=width, num_frames=num_frames, warmup=False, **kwargs)
if self.nvtx_profile:
cudart.cudaProfilerStop()
demo/Diffusion/demo_diffusion/utils_modelopt.py
Executable
+853
View File
@@ -0,0 +1,853 @@
#
# 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
from collections import defaultdict
from random import choice, shuffle
from typing import Set
import modelopt.torch.quantization as mtq
import numpy as np
import onnx_graphsurgeon as gs
import torch
import torch.nn.functional as F
from diffusers.models.attention_processor import (
Attention,
AttnProcessor,
FluxAttnProcessor2_0,
)
from diffusers.models.lora import LoRACompatibleConv, LoRACompatibleLinear
from modelopt.torch.quantization import utils as quant_utils
from modelopt.torch.quantization.calib.max import MaxCalibrator
from PIL import Image
from torch.utils.data import Dataset, Sampler
import onnx
USE_PEFT = True
try:
from peft.tuners.lora.layer import Conv2d as PEFTLoRAConv2d
from peft.tuners.lora.layer import Linear as PEFTLoRALinear
except ModuleNotFoundError:
USE_PEFT = False
class PercentileCalibrator(MaxCalibrator):
def __init__(self, num_bits=8, axis=None, unsigned=False, track_amax=False, **kwargs):
super().__init__(num_bits, axis, unsigned, track_amax)
self.percentile = kwargs["percentile"]
self.total_step = kwargs["total_step"]
self.collect_method = kwargs["collect_method"]
self.data = {}
self.i = 0
def collect(self, x):
"""Tracks the absolute max of all tensors.
Args:
x: A tensor
Raises:
RuntimeError: If amax shape changes
"""
# Swap axis to reduce.
axis = self._axis if isinstance(self._axis, (list, tuple)) else [self._axis]
# Handle negative axis.
axis = [x.dim() + i if isinstance(i, int) and i < 0 else i for i in axis]
reduce_axis = []
for i in range(x.dim()):
if i not in axis:
reduce_axis.append(i)
local_amax = quant_utils.reduce_amax(x, axis=reduce_axis).detach()
_cur_step = self.i % self.total_step
if _cur_step not in self.data.keys():
self.data[_cur_step] = local_amax
else:
if self.collect_method == "global_min":
self.data[_cur_step] = torch.min(self.data[_cur_step], local_amax)
elif self.collect_method == "min-max" or self.collect_method == "mean-max":
self.data[_cur_step] = torch.max(self.data[_cur_step], local_amax)
else:
self.data[_cur_step] += local_amax
if self._track_amax:
raise NotImplementedError
self.i += 1
def compute_amax(self):
"""Return the absolute max of all tensors collected."""
up_lim = int(self.total_step * self.percentile)
if self.collect_method == "min-mean":
amaxs_values = [self.data[i] / self.total_step for i in range(0, up_lim)]
else:
amaxs_values = [self.data[i] for i in range(0, up_lim)]
if self.collect_method == "mean-max":
act_amax = torch.vstack(amaxs_values).mean(axis=0)[0]
else:
act_amax = torch.vstack(amaxs_values).min(axis=0)[0]
self._calib_amax = act_amax
return self._calib_amax
def __str__(self):
s = "PercentileCalibrator"
return s.format(**self.__dict__)
def __repr__(self):
s = "PercentileCalibrator("
s += super(MaxCalibrator, self).__repr__()
s += " calib_amax={_calib_amax}"
if self._track_amax:
s += " amaxs={_amaxs}"
s += ")"
return s.format(**self.__dict__)
def filter_func(name):
pattern = re.compile(
r".*(time_emb_proj|time_embedding|conv_in|conv_out|conv_shortcut|add_embedding|pos_embed|time_text_embed|context_embedder|norm_out|proj_out).*"
)
return pattern.match(name) is not None
def filter_func_no_proj_out(name): # used for Flux
pattern = re.compile(
r".*(time_emb_proj|time_embedding|conv_in|conv_out|conv_shortcut|add_embedding|pos_embed|time_text_embed|context_embedder|norm_out|x_embedder).*"
)
return pattern.match(name) is not None
def quantize_lvl(model_id, backbone, quant_level=2.5, linear_only=False, enable_conv_3d=True):
"""
We should disable the unwanted quantizer when exporting the onnx
Because in the current modelopt setting, it will load the quantizer amax for all the layers even
if we didn't add that unwanted layer into the config during the calibration
"""
for name, module in backbone.named_modules():
if isinstance(module, torch.nn.Conv2d):
if linear_only:
module.input_quantizer.disable()
module.weight_quantizer.disable()
else:
module.input_quantizer.enable()
module.weight_quantizer.enable()
elif isinstance(module, torch.nn.Linear):
if (
(quant_level >= 2 and "ff.net" in name)
or (quant_level >= 2.5 and ("to_q" in name or "to_k" in name or "to_v" in name))
or quant_level >= 3
) and name != "proj_out": # Disable the final output layer from flux model
module.input_quantizer.enable()
module.weight_quantizer.enable()
else:
module.input_quantizer.disable()
module.weight_quantizer.disable()
elif isinstance(module, torch.nn.Conv3d) and not enable_conv_3d:
"""
Error: Torch bug, ONNX export failed due to unknown kernel shape in QuantConv3d.
TRT_FP8QuantizeLinear and TRT_FP8DequantizeLinear operations in UNetSpatioTemporalConditionModel for svd
cause issues. Inputs on different devices (CUDA vs CPU) may contribute to the problem.
"""
module.input_quantizer.disable()
module.weight_quantizer.disable()
elif isinstance(module, Attention):
# TRT only supports FP8 MHA with head_size % 16 == 0.
head_size = int(module.inner_dim / module.heads)
if quant_level >= 4 and head_size % 16 == 0:
module.q_bmm_quantizer.enable()
module.k_bmm_quantizer.enable()
module.v_bmm_quantizer.enable()
module.softmax_quantizer.enable()
if model_id.startswith("flux.1"):
if name.startswith("transformer_blocks"):
module.bmm2_output_quantizer.enable()
else:
module.bmm2_output_quantizer.disable()
setattr(module, "_disable_fp8_mha", False)
else:
module.q_bmm_quantizer.disable()
module.k_bmm_quantizer.disable()
module.v_bmm_quantizer.disable()
module.softmax_quantizer.disable()
module.bmm2_output_quantizer.disable()
setattr(module, "_disable_fp8_mha", True)
def fp8_mha_disable(backbone, quantized_mha_output: bool = True):
def mha_filter_func(name):
pattern = re.compile(
r".*(q_bmm_quantizer|k_bmm_quantizer|v_bmm_quantizer|softmax_quantizer).*"
if quantized_mha_output
else r".*(q_bmm_quantizer|k_bmm_quantizer|v_bmm_quantizer|softmax_quantizer|bmm2_output_quantizer).*"
)
return pattern.match(name) is not None
if hasattr(F, "scaled_dot_product_attention"):
mtq.disable_quantizer(backbone, mha_filter_func)
def get_int8_config(
model,
quant_level=3,
alpha=0.8,
percentile=1.0,
num_inference_steps=20,
collect_method="min-mean",
):
quant_config = {
"quant_cfg": {
"*lm_head*": {"enable": False},
"*output_layer*": {"enable": False},
"*output_quantizer": {"enable": False},
"default": {"num_bits": 8, "axis": None},
},
"algorithm": {"method": "smoothquant", "alpha": alpha},
}
for name, module in model.named_modules():
w_name = f"{name}*weight_quantizer"
i_name = f"{name}*input_quantizer"
if w_name in quant_config["quant_cfg"].keys() or i_name in quant_config["quant_cfg"].keys():
continue
if filter_func(name):
continue
if isinstance(module, (torch.nn.Linear, LoRACompatibleLinear)):
if (
(quant_level >= 2 and "ff.net" in name)
or (quant_level >= 2.5 and ("to_q" in name or "to_k" in name or "to_v" in name))
or quant_level == 3
):
quant_config["quant_cfg"][w_name] = {"num_bits": 8, "axis": 0}
quant_config["quant_cfg"][i_name] = {"num_bits": 8, "axis": -1}
elif isinstance(module, (torch.nn.Conv2d, LoRACompatibleConv)):
quant_config["quant_cfg"][w_name] = {"num_bits": 8, "axis": 0}
quant_config["quant_cfg"][i_name] = {
"num_bits": 8,
"axis": None,
"calibrator": (
PercentileCalibrator,
(),
{
"num_bits": 8,
"axis": None,
"percentile": percentile,
"total_step": num_inference_steps,
"collect_method": collect_method,
},
),
}
return quant_config
SD_FP8_FP16_DEFAULT_CONFIG = {
"quant_cfg": {
"*weight_quantizer": {"num_bits": (4, 3), "axis": None, "trt_high_precision_dtype": "Half"},
"*input_quantizer": {"num_bits": (4, 3), "axis": None, "trt_high_precision_dtype": "Half"},
"*output_quantizer": {"enable": False},
"*q_bmm_quantizer": {"num_bits": (4, 3), "axis": None, "trt_high_precision_dtype": "Half"},
"*k_bmm_quantizer": {"num_bits": (4, 3), "axis": None, "trt_high_precision_dtype": "Half"},
"*v_bmm_quantizer": {"num_bits": (4, 3), "axis": None, "trt_high_precision_dtype": "Half"},
"*softmax_quantizer": {
"num_bits": (4, 3),
"axis": None,
"trt_high_precision_dtype": "Half",
},
"default": {"enable": False},
},
"algorithm": "max",
}
SD_FP8_BF16_DEFAULT_CONFIG = {
"quant_cfg": {
"*weight_quantizer": {"num_bits": (4, 3), "axis": None, "trt_high_precision_dtype": "BFloat16"},
"*input_quantizer": {"num_bits": (4, 3), "axis": None, "trt_high_precision_dtype": "BFloat16"},
"*output_quantizer": {"enable": False},
"*q_bmm_quantizer": {"num_bits": (4, 3), "axis": None, "trt_high_precision_dtype": "BFloat16"},
"*k_bmm_quantizer": {"num_bits": (4, 3), "axis": None, "trt_high_precision_dtype": "BFloat16"},
"*v_bmm_quantizer": {"num_bits": (4, 3), "axis": None, "trt_high_precision_dtype": "BFloat16"},
"*softmax_quantizer": {
"num_bits": (4, 3),
"axis": None,
"trt_high_precision_dtype": "BFloat16",
},
"default": {"enable": False},
},
"algorithm": "max",
}
SD_FP8_BF16_FLUX_MMDIT_BMM2_FP8_OUTPUT_CONFIG = {
"quant_cfg": {
"*weight_quantizer": {"num_bits": (4, 3), "axis": None, "trt_high_precision_dtype": "BFloat16"},
"*input_quantizer": {"num_bits": (4, 3), "axis": None, "trt_high_precision_dtype": "BFloat16"},
"*output_quantizer": {"enable": False},
"*q_bmm_quantizer": {"num_bits": (4, 3), "axis": None, "trt_high_precision_dtype": "BFloat16"},
"*k_bmm_quantizer": {"num_bits": (4, 3), "axis": None, "trt_high_precision_dtype": "BFloat16"},
"*v_bmm_quantizer": {"num_bits": (4, 3), "axis": None, "trt_high_precision_dtype": "BFloat16"},
"*softmax_quantizer": {
"num_bits": (4, 3),
"axis": None,
"trt_high_precision_dtype": "BFloat16",
},
"transformer_blocks*bmm2_output_quantizer": {
"num_bits": (4, 3),
"axis": None,
"trt_high_precision_dtype": "BFloat16",
},
"default": {"enable": False},
},
"algorithm": "max",
}
SD_FP8_FP32_DEFAULT_CONFIG = {
"quant_cfg": {
"*weight_quantizer": {"num_bits": (4, 3), "axis": None, "trt_high_precision_dtype": "Float"},
"*input_quantizer": {"num_bits": (4, 3), "axis": None, "trt_high_precision_dtype": "Float"},
"*output_quantizer": {"enable": False},
"*q_bmm_quantizer": {"num_bits": (4, 3), "axis": None, "trt_high_precision_dtype": "Float"},
"*k_bmm_quantizer": {"num_bits": (4, 3), "axis": None, "trt_high_precision_dtype": "Float"},
"*v_bmm_quantizer": {"num_bits": (4, 3), "axis": None, "trt_high_precision_dtype": "Float"},
"*softmax_quantizer": {
"num_bits": (4, 3),
"axis": None,
"trt_high_precision_dtype": "Float",
},
"default": {"enable": False},
},
"algorithm": "max",
}
def set_fmha(denoiser, is_flux=False):
for name, module in denoiser.named_modules():
if isinstance(module, Attention):
if is_flux:
module.set_processor(FluxAttnProcessor2_0())
else:
module.set_processor(AttnProcessor())
def check_lora(model):
for name, module in model.named_modules():
if isinstance(module, (LoRACompatibleConv, LoRACompatibleLinear)):
assert (
module.lora_layer is None
), f"To quantize {name}, LoRA layer should be fused/merged. Please fuse the LoRA layer before quantization."
elif USE_PEFT and isinstance(module, (PEFTLoRAConv2d, PEFTLoRALinear)):
assert (
module.merged
), f"To quantize {name}, LoRA layer should be fused/merged. Please fuse the LoRA layer before quantization."
def generate_fp8_scales(unet):
# temporary solution due to a known bug in torch.onnx._dynamo_export
for _, module in unet.named_modules():
if isinstance(module, (torch.nn.Linear, torch.nn.Conv2d)) and (
hasattr(module.input_quantizer, "_amax") and module.input_quantizer is not None
):
module.input_quantizer._num_bits = 8
module.weight_quantizer._num_bits = 8
module.input_quantizer._amax = module.input_quantizer._amax * (127 / 448.0)
module.weight_quantizer._amax = module.weight_quantizer._amax * (127 / 448.0)
elif isinstance(module, Attention) and (
hasattr(module.q_bmm_quantizer, "_amax")
and module.q_bmm_quantizer is not None
and hasattr(module.k_bmm_quantizer, "_amax")
and module.k_bmm_quantizer is not None
and hasattr(module.v_bmm_quantizer, "_amax")
and module.v_bmm_quantizer is not None
and hasattr(module.softmax_quantizer, "_amax")
and module.softmax_quantizer is not None
):
module.q_bmm_quantizer._num_bits = 8
module.q_bmm_quantizer._amax = module.q_bmm_quantizer._amax * (127 / 448.0)
module.k_bmm_quantizer._num_bits = 8
module.k_bmm_quantizer._amax = module.k_bmm_quantizer._amax * (127 / 448.0)
module.v_bmm_quantizer._num_bits = 8
module.v_bmm_quantizer._amax = module.v_bmm_quantizer._amax * (127 / 448.0)
module.softmax_quantizer._num_bits = 8
module.softmax_quantizer._amax = module.softmax_quantizer._amax * (127 / 448.0)
def get_parent_nodes(node):
"""
Returns list of input producer nodes for the given node.
"""
parents = []
for tensor in node.inputs:
# If the tensor is not a constant or graph input and has a producer,
# the producer is a parent of node `node`
if len(tensor.inputs) == 1:
parents.append(tensor.inputs[0])
return parents
def get_child_nodes(node):
"""
Returns list of output consumer nodes for the given node.
"""
children = []
for tensor in node.outputs:
for consumer in tensor.outputs: # Traverse all consumer of the tensor
children.append(consumer)
return children
def has_path_type(node, graph, path_type, is_forward, wild_card_types, path_nodes):
"""
Return pattern nodes for the given path_type.
"""
if not path_type:
# All types matched
return True
# Check if current non-wild node type does not match the expected path type
node_type = node.op
is_match = node_type == path_type[0]
is_wild_match = node_type in wild_card_types
if not is_match and not is_wild_match:
return False
if is_match:
path_nodes.append(node)
next_path_type = path_type[1:]
else:
next_path_type = path_type[:]
if is_forward:
next_level_nodes = get_child_nodes(node)
else:
next_level_nodes = get_parent_nodes(node)
# Check if any child (forward path) or parent (backward path) can match the remaining path types
for next_node in next_level_nodes:
sub_path = []
if has_path_type(next_node, graph, next_path_type, is_forward, wild_card_types, sub_path):
path_nodes.extend(sub_path)
return True
# Path type matches if there is no remaining types to match
return not next_path_type
def insert_cast(graph, input_tensor, attrs):
"""
Create a cast layer using tensor as input.
"""
output_tensor = gs.Variable(name=f"{input_tensor.name}/Cast_output", dtype=attrs["to"])
next_node_list = input_tensor.outputs.copy()
graph.layer(
op="Cast",
name=f"{input_tensor.name}/Cast",
inputs=[input_tensor],
outputs=[output_tensor],
attrs=attrs,
)
# use cast output as input to next node
for next_node in next_node_list:
for idx, next_input in enumerate(next_node.inputs):
if next_input.name == input_tensor.name:
next_node.inputs[idx] = output_tensor
def cast_layernorm_io(graph):
"""
Cast LayerNormalization scale and bias inputs from FP16 to FP32.
In INT8 quantized graphs, DequantizeLinear outputs Float32 activations,
but LayerNorm scale/bias remain FP16 from the original model, causing
a type mismatch with --strongly-typed TensorRT builds.
"""
layernorm_nodes = [node for node in graph.nodes if node.op == "LayerNormalization"]
print(f"Found {len(layernorm_nodes)} LayerNormalization nodes to fix")
for node in layernorm_nodes:
# LayerNormalization inputs: 0=X (data), 1=Scale, 2=B (bias, optional)
for i in range(1, len(node.inputs)):
input_tensor = node.inputs[i]
if input_tensor.name and hasattr(input_tensor, 'dtype') and input_tensor.dtype == np.float16:
insert_cast(graph, input_tensor=input_tensor, attrs={"to": np.float32})
def cast_convtranspose_io(graph):
"""
Fix ConvTranspose input/output type mismatches for strongly-typed TRT builds.
In mixed-precision graphs (e.g. BF16 Stable Cascade VQGAN), architectural FP16->FP32
casts can leave a ConvTranspose with a FP32 activation input but FP16 kernel weights.
We cast the activation to match the kernel dtype, then cast the output back to the
original activation dtype so surrounding FP32 ops (e.g. residual Add) are unaffected.
"""
convtranspose_nodes = [node for node in graph.nodes if node.op == "ConvTranspose"]
fixed = 0
for node in convtranspose_nodes:
if len(node.inputs) < 2:
continue
act_input = node.inputs[0]
kernel = node.inputs[1]
if act_input.dtype is None or kernel.dtype is None or act_input.dtype == kernel.dtype:
continue
orig_dtype = act_input.dtype # e.g. np.dtype('float32')
target_dtype = kernel.dtype.type # e.g. np.float16
insert_cast(graph, input_tensor=act_input, attrs={"to": target_dtype})
# Update the output dtype to match and cast back, so downstream FP32 ops are unaffected.
for out in node.outputs:
if out.name and out.dtype == orig_dtype:
out.dtype = target_dtype
insert_cast(graph, input_tensor=out, attrs={"to": orig_dtype.type})
fixed += 1
print(f"Fixed {fixed} ConvTranspose input/output type mismatches")
def convert_zp_fp8(onnx_graph):
"""
Convert Q/DQ zero datatype from INT8 to FP8.
"""
# Find all zero constant nodes
qdq_zero_nodes = set()
for node in onnx_graph.graph.node:
if node.op_type == "QuantizeLinear":
if len(node.input) > 2:
qdq_zero_nodes.add(node.input[2])
print(f"Found {len(qdq_zero_nodes)} QDQ pairs")
# Convert zero point datatype from INT8 to FP8.
for node in onnx_graph.graph.node:
if node.output[0] in qdq_zero_nodes:
node.attribute[0].t.data_type = onnx.TensorProto.FLOAT8E4M3FN
return onnx_graph
def cast_resize_io(graph, output_dtype=np.float16):
"""
Add cast nodes to Resize nodes I/O because Resize needs to be run in fp32.
Inputs are cast to FP32, outputs are cast back to output_dtype (FP16 or BF16).
"""
resize_nodes = [node for node in graph.nodes if node.op == "Resize"]
print(f"Found {len(resize_nodes)} Resize nodes to fix")
for resize_node in resize_nodes:
# Skip Resize nodes whose data input is already FP32 — no casting needed.
if resize_node.inputs[0].dtype == np.float32:
continue
for i, input_tensor in enumerate(resize_node.inputs):
SIZES_INPUT_INDEX = 3 # Optional input "sizes" at index 3 must be in INT64. Skip cast for this input.
if i != SIZES_INPUT_INDEX and input_tensor.name:
insert_cast(graph, input_tensor=input_tensor, attrs={"to": np.float32})
for output_tensor in resize_node.outputs:
if output_tensor.name:
insert_cast(graph, input_tensor=output_tensor, attrs={"to": output_dtype})
def cast_fp8_mha_io(graph):
r"""
Insert three cast ops.
The first cast will be added before the input0 of MatMul to cast fp16 to fp32.
The second cast will be added before the input1 of MatMul to cast fp16 to fp32.
The third cast will be added after the output of MatMul to cast fp32 back to fp16.
Q Q
| |
DQ DQ
| |
Cast Cast
(fp16 to fp32) (fp16 to fp32)
\ /
\ /
\ /
MatMul
|
Cast (fp32 to fp16)
|
Q
|
DQ
The insertion of Cast ops in the FP8 MHA part actually forbids the MHAs to run
with FP16 accumulation because TensorRT only has FP32 accumulation kernels for FP8 MHAs.
"""
# Find FP8 MHA pattern.
# Match FP8 MHA: Q -> DQ -> BMM1 -> (Mul/Div) -> (Add) -> Softmax -> (Cast) -> Q -> DQ -> BMM2 -> Q -> DQ
softmax_bmm1_chain_type = ["Softmax", "MatMul", "DequantizeLinear", "QuantizeLinear"]
softmax_bmm2_chain_type = [
"Softmax",
"QuantizeLinear",
"DequantizeLinear",
"MatMul",
"QuantizeLinear",
"DequantizeLinear",
]
wild_card_types = [
"Div",
"Mul",
"ConstMul",
"Add",
"BiasAdd",
"Reshape",
"Transpose",
"Flatten",
"Cast",
]
fp8_mha_partitions = []
for node in graph.nodes:
if node.op == "Softmax":
fp8_mha_partition = []
if has_path_type(
node, graph, softmax_bmm1_chain_type, False, wild_card_types, fp8_mha_partition
) and has_path_type(
node, graph, softmax_bmm2_chain_type, True, wild_card_types, fp8_mha_partition
):
if (
len(fp8_mha_partition) == 10
and fp8_mha_partition[1].op == "MatMul"
and fp8_mha_partition[7].op == "MatMul"
):
fp8_mha_partitions.append(fp8_mha_partition)
print(f"Found {len(fp8_mha_partitions)} FP8 attentions")
# Insert Cast nodes for BMM1 and BMM2.
for fp8_mha_partition in fp8_mha_partitions:
bmm1_node = fp8_mha_partition[1]
insert_cast(graph, input_tensor=bmm1_node.inputs[0], attrs={"to": np.float32})
insert_cast(graph, input_tensor=bmm1_node.inputs[1], attrs={"to": np.float32})
insert_cast(graph, input_tensor=bmm1_node.outputs[0], attrs={"to": np.float16})
bmm2_node = fp8_mha_partition[7]
insert_cast(graph, input_tensor=bmm2_node.inputs[0], attrs={"to": np.float32})
insert_cast(graph, input_tensor=bmm2_node.inputs[1], attrs={"to": np.float32})
insert_cast(graph, input_tensor=bmm2_node.outputs[0], attrs={"to": np.float16})
def set_quant_precision(quant_config, precision: str = "Half"):
for key in quant_config["quant_cfg"]:
if "trt_high_precision_dtype" in quant_config["quant_cfg"][key]:
quant_config["quant_cfg"][key]["trt_high_precision_dtype"] = precision
def convert_fp16_io(graph):
"""
Convert graph I/O to FP16.
"""
for input_tensor in graph.inputs:
input_tensor.dtype = onnx.TensorProto.FLOAT16
for output_tensor in graph.outputs:
output_tensor.dtype = onnx.TensorProto.FLOAT16
def random_resize(cur_size: int):
"""
Randomly selects a new resolution for an image based on its current aspect ratio.
This function determines the current aspect ratio of an image, selects a new aspect ratio
from predefined choices depending on whether the current aspect ratio is square,
portrait, or landscape, and returns the corresponding resolution from a provided mapping.
Parameters:
cur_size (int): A tuple (width, height) representing the current resolution of the image.
resolution_to_aspects (dict[float, tuple[int, int]]): A mapping of aspect ratios (floats)
to their corresponding resolutions as tuples of (width, height).
Returns:
tuple[int, int]: A tuple (new_width, new_height) representing the newly selected resolution.
Raises:
KeyError: If the chosen aspect ratio is not present in the `resolution_to_aspects` dictionary.
Notes:
- For square images (aspect ratio = 1), the function selects from aspect ratios 1.25, 0.8, 1.5, and 0.667.
- For landscape images (aspect ratio > 1), the function selects from aspect ratios 1.778, 1.25, and 1.5.
- For portrait images (aspect ratio < 1), the function selects from aspect ratios 0.563, 0.8, and 0.667.
"""
resolution_to_aspects = {
1.0: (1024, 1024),
1.778: (768, 1344),
0.563: (1344, 768),
1.25: (896, 1152),
0.8: (1152, 896),
1.5: (832, 1216),
0.667: (1216, 832),
}
cur_aspect_ratio = round(cur_size[1] / cur_size[0], 3)
if cur_aspect_ratio == 1:
new_aspect_ratio = choice((1.25, 0.8, 1.5, 0.667))
new_res = resolution_to_aspects[new_aspect_ratio]
elif cur_aspect_ratio > 1:
new_aspect_ratio = choice((1.778, 1.25, 1.5))
new_res = resolution_to_aspects[new_aspect_ratio]
else:
# cur_aspect_ratio < 1
new_aspect_ratio = choice((0.563, 0.8, 0.667))
new_res = resolution_to_aspects[new_aspect_ratio]
return new_res
class PromptImageDataset(Dataset):
def __init__(
self,
root_dir,
):
"""
Args:
root_dir (str): Directory with all the images and the prompt file.
"""
self.root_dir = root_dir
self.possible_resolutions = {1024, 768, 1344, 896, 832, 1216}
self.global_idx_template = "{} | {} | {}"
self.prompts_by_size = defaultdict(list)
self.images_by_size = defaultdict(list)
self.images = []
self.prompts = []
self.images_size = []
# self.global_idx_2_group = dict()
# self.global_idx_to_group_idx = dict()
self.group_to_global_idx = {}
for idx, file in enumerate(os.listdir(os.path.join(self.root_dir, "prompts"))):
if not file.endswith(".txt"):
continue
file_name = os.path.splitext(file)[0]
image_path = os.path.join(
self.root_dir,
"inputs",
f"{file_name}.png",
)
with Image.open(image_path) as img, open(os.path.join(self.root_dir, "prompts", file), "r") as f:
prompt = "\n".join(f.readlines())
std_img_size = (
self.closest_value(img.size[0], self.possible_resolutions),
self.closest_value(img.size[1], self.possible_resolutions),
)
self.images_by_size[std_img_size].append(image_path)
self.prompts_by_size[std_img_size].append(prompt)
self.images.append(image_path)
self.prompts.append(prompt)
self.images_size.append(std_img_size)
# create a unique key that map group and index inside the group to a global index
in_group_idx = len(self.images_by_size[std_img_size]) - 1
group_idx_key = self.global_idx_template.format(std_img_size[0], std_img_size[1], in_group_idx)
self.group_to_global_idx[group_idx_key] = len(self.images) - 1
assert len(self.images) == len(self.prompts)
assert len(self.images) == len(self.group_to_global_idx)
@staticmethod
def closest_value(target: int, candidates: Set[int]):
"""
Find the closest value to the target from a set of candidate values.
Args:
target (int): The integer to compare against.
candidates (set): A set of integers as candidates.
Returns:
int: The closest value from the candidates.
"""
if not candidates:
raise ValueError("The candidates set cannot be empty.")
# Use the min function with a key that computes the absolute difference
return min(candidates, key=lambda x: abs(x - target))
def __len__(self):
return len(self.images)
def __getitem__(self, idx):
"""
Returns:
image (Tensor): Transformed image.
prompt (str): Corresponding text prompt.
"""
if torch.is_tensor(idx):
idx = idx.tolist()
prompt = self.prompts[idx]
image = self.images[idx]
image_size = self.images_size[idx]
return image, prompt, image_size
class SameSizeSampler(Sampler):
def __init__(self, dataset: PromptImageDataset, batch_size: int):
"""
Custom sampler that creates batches of images with the same size
Args:
dataset (SameSizeImageDataset): Dataset to sample from
batch_size (int): Number of images per batch
"""
super().__init__(dataset)
self.dataset = dataset
self.batch_size = batch_size
# Prepare size groups with indices
self.size_groups = {}
for size, image_paths in self.dataset.images_by_size.items():
# Create a list of indices for this size group
self.size_groups[size] = list(range(len(image_paths)))
def __iter__(self):
"""
Iteration method that yields indices for batches of same-size images
"""
# Create a copy of size groups to shuffle
size_groups_copy = {std_img_size: indices.copy() for std_img_size, indices in self.size_groups.items()}
# Shuffle each size group
for std_img_size, indices in size_groups_copy.items():
shuffle(indices)
# Iterate through size groups
for std_img_size, indices in size_groups_copy.items():
# Batch indices of the same size
for i in range(0, len(indices), self.batch_size):
# Yield batch indices for this size
batch_group_idxs = indices[i : min(i + self.batch_size, len(indices))]
for in_group_idx in batch_group_idxs:
group_idx_key = self.dataset.global_idx_template.format(
std_img_size[0], std_img_size[1], in_group_idx
)
batch_global_idx = self.dataset.group_to_global_idx[group_idx_key]
# batch_global_idxs.append(batch_global_idx)
yield batch_global_idx
def __len__(self):
"""
Total number of batches
"""
return len(self.dataset.images) // self.batch_size
def custom_collate(data):
"""
Custom collate function to handle batches of same-size images
Args:
dataset (SameSizeImageDataset): Dataset instance
batch (list): List of global indices
Returns:
tuple: Batched images and their size
"""
# Group images by their size
images, prompts, image_sizes = tuple(map(list, zip(*data)))
assert len(images) > 0
new_img_size = random_resize(image_sizes[0])
batch_images = []
for image in images:
with Image.open(image) as image:
image = image.convert("RGB").resize(size=new_img_size, resample=Image.LANCZOS)
image = np.array(image)
image = np.transpose(image, axes=(-1, 0, 1))
image = torch.from_numpy(image).float() / 127.5 - 1.0
batch_images.append(image)
batch_images = torch.stack(batch_images, dim=0)
return batch_images, prompts
def infinite_dataloader(dataloader):
while True:
for batch in dataloader:
yield batch
demo/Diffusion/demo_diffusion/utils_sd3/__init__.py
View File
demo/Diffusion/demo_diffusion/utils_sd3/mmdit.py
+641
View File
@@ -0,0 +1,641 @@
# MIT License
# Copyright (c) 2024 Stability AI
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
# The above copyright notice and this permission notice shall be included in all
# copies or substantial portions of the Software.
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.
import math
from typing import Dict, Optional
import numpy as np
import torch
import torch.nn as nn
from einops import rearrange, repeat
from demo_diffusion.utils_sd3.other_impls import Mlp, attention
class PatchEmbed(nn.Module):
""" 2D Image to Patch Embedding"""
def __init__(
self,
img_size: Optional[int] = 224,
patch_size: int = 16,
in_chans: int = 3,
embed_dim: int = 768,
flatten: bool = True,
bias: bool = True,
strict_img_size: bool = True,
dynamic_img_pad: bool = False,
dtype=None,
device=None,
):
super().__init__()
self.patch_size = (patch_size, patch_size)
if img_size is not None:
self.img_size = (img_size, img_size)
self.grid_size = tuple([s // p for s, p in zip(self.img_size, self.patch_size)])
self.num_patches = self.grid_size[0] * self.grid_size[1]
else:
self.img_size = None
self.grid_size = None
self.num_patches = None
# flatten spatial dim and transpose to channels last, kept for bwd compat
self.flatten = flatten
self.strict_img_size = strict_img_size
self.dynamic_img_pad = dynamic_img_pad
self.proj = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_size, stride=patch_size, bias=bias, dtype=dtype, device=device)
def forward(self, x):
B, C, H, W = x.shape
x = self.proj(x)
if self.flatten:
x = x.flatten(2).transpose(1, 2) # NCHW -> NLC
return x
def modulate(x, shift, scale):
if shift is None:
shift = torch.zeros_like(scale)
return x * (1 + scale.unsqueeze(1)) + shift.unsqueeze(1)
#################################################################################
# Sine/Cosine Positional Embedding Functions #
#################################################################################
def get_2d_sincos_pos_embed(embed_dim, grid_size, cls_token=False, extra_tokens=0, scaling_factor=None, offset=None):
"""
grid_size: int of the grid height and width
return:
pos_embed: [grid_size*grid_size, embed_dim] or [1+grid_size*grid_size, embed_dim] (w/ or w/o cls_token)
"""
grid_h = np.arange(grid_size, dtype=np.float32)
grid_w = np.arange(grid_size, dtype=np.float32)
grid = np.meshgrid(grid_w, grid_h) # here w goes first
grid = np.stack(grid, axis=0)
if scaling_factor is not None:
grid = grid / scaling_factor
if offset is not None:
grid = grid - offset
grid = grid.reshape([2, 1, grid_size, grid_size])
pos_embed = get_2d_sincos_pos_embed_from_grid(embed_dim, grid)
if cls_token and extra_tokens > 0:
pos_embed = np.concatenate([np.zeros([extra_tokens, embed_dim]), pos_embed], axis=0)
return pos_embed
def get_2d_sincos_pos_embed_from_grid(embed_dim, grid):
assert embed_dim % 2 == 0
# use half of dimensions to encode grid_h
emb_h = get_1d_sincos_pos_embed_from_grid(embed_dim // 2, grid[0]) # (H*W, D/2)
emb_w = get_1d_sincos_pos_embed_from_grid(embed_dim // 2, grid[1]) # (H*W, D/2)
emb = np.concatenate([emb_h, emb_w], axis=1) # (H*W, D)
return emb
def get_1d_sincos_pos_embed_from_grid(embed_dim, pos):
"""
embed_dim: output dimension for each position
pos: a list of positions to be encoded: size (M,)
out: (M, D)
"""
assert embed_dim % 2 == 0
omega = np.arange(embed_dim // 2, dtype=np.float64)
omega /= embed_dim / 2.0
omega = 1.0 / 10000**omega # (D/2,)
pos = pos.reshape(-1) # (M,)
out = np.einsum("m,d->md", pos, omega) # (M, D/2), outer product
emb_sin = np.sin(out) # (M, D/2)
emb_cos = np.cos(out) # (M, D/2)
return np.concatenate([emb_sin, emb_cos], axis=1) # (M, D)
#################################################################################
# Embedding Layers for Timesteps and Class Labels #
#################################################################################
class TimestepEmbedder(nn.Module):
"""Embeds scalar timesteps into vector representations."""
def __init__(self, hidden_size, frequency_embedding_size=256, dtype=None, device=None):
super().__init__()
self.mlp = nn.Sequential(
nn.Linear(frequency_embedding_size, hidden_size, bias=True, dtype=dtype, device=device),
nn.SiLU(),
nn.Linear(hidden_size, hidden_size, bias=True, dtype=dtype, device=device),
)
self.frequency_embedding_size = frequency_embedding_size
@staticmethod
def timestep_embedding(t, dim, max_period=10000):
"""
Create sinusoidal timestep embeddings.
:param t: a 1-D Tensor of N indices, one per batch element.
These may be fractional.
:param dim: the dimension of the output.
:param max_period: controls the minimum frequency of the embeddings.
:return: an (N, D) Tensor of positional embeddings.
"""
half = dim // 2
freqs = torch.exp(
-math.log(max_period)
* torch.arange(start=0, end=half, dtype=torch.float32)
/ half
).to(device=t.device)
args = t[:, None].float() * freqs[None]
embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1)
if dim % 2:
embedding = torch.cat([embedding, torch.zeros_like(embedding[:, :1])], dim=-1)
if torch.is_floating_point(t):
embedding = embedding.to(dtype=t.dtype)
return embedding
def forward(self, t, dtype, **kwargs):
t_freq = self.timestep_embedding(t, self.frequency_embedding_size).to(dtype)
t_emb = self.mlp(t_freq)
return t_emb
class VectorEmbedder(nn.Module):
"""Embeds a flat vector of dimension input_dim"""
def __init__(self, input_dim: int, hidden_size: int, dtype=None, device=None):
super().__init__()
self.mlp = nn.Sequential(
nn.Linear(input_dim, hidden_size, bias=True, dtype=dtype, device=device),
nn.SiLU(),
nn.Linear(hidden_size, hidden_size, bias=True, dtype=dtype, device=device),
)
def forward(self, x: torch.Tensor) -> torch.Tensor:
return self.mlp(x)
#################################################################################
# Core DiT Model #
#################################################################################
def split_qkv(qkv, head_dim):
qkv = qkv.reshape(qkv.shape[0], qkv.shape[1], 3, -1, head_dim).movedim(2, 0)
return qkv[0], qkv[1], qkv[2]
def optimized_attention(qkv, num_heads):
return attention(qkv[0], qkv[1], qkv[2], num_heads)
class SelfAttention(nn.Module):
ATTENTION_MODES = ("xformers", "torch", "torch-hb", "math", "debug")
def __init__(
self,
dim: int,
num_heads: int = 8,
qkv_bias: bool = False,
qk_scale: Optional[float] = None,
attn_mode: str = "xformers",
pre_only: bool = False,
qk_norm: Optional[str] = None,
rmsnorm: bool = False,
dtype=None,
device=None,
):
super().__init__()
self.num_heads = num_heads
self.head_dim = dim // num_heads
self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias, dtype=dtype, device=device)
if not pre_only:
self.proj = nn.Linear(dim, dim, dtype=dtype, device=device)
assert attn_mode in self.ATTENTION_MODES
self.attn_mode = attn_mode
self.pre_only = pre_only
if qk_norm == "rms":
self.ln_q = RMSNorm(self.head_dim, elementwise_affine=True, eps=1.0e-6, dtype=dtype, device=device)
self.ln_k = RMSNorm(self.head_dim, elementwise_affine=True, eps=1.0e-6, dtype=dtype, device=device)
elif qk_norm == "ln":
self.ln_q = nn.LayerNorm(self.head_dim, elementwise_affine=True, eps=1.0e-6, dtype=dtype, device=device)
self.ln_k = nn.LayerNorm(self.head_dim, elementwise_affine=True, eps=1.0e-6, dtype=dtype, device=device)
elif qk_norm is None:
self.ln_q = nn.Identity()
self.ln_k = nn.Identity()
else:
raise ValueError(qk_norm)
def pre_attention(self, x: torch.Tensor):
B, L, C = x.shape
qkv = self.qkv(x)
q, k, v = split_qkv(qkv, self.head_dim)
q = self.ln_q(q).reshape(q.shape[0], q.shape[1], -1)
k = self.ln_k(k).reshape(q.shape[0], q.shape[1], -1)
return (q, k, v)
def post_attention(self, x: torch.Tensor) -> torch.Tensor:
assert not self.pre_only
x = self.proj(x)
return x
def forward(self, x: torch.Tensor) -> torch.Tensor:
(q, k, v) = self.pre_attention(x)
x = attention(q, k, v, self.num_heads)
x = self.post_attention(x)
return x
class RMSNorm(torch.nn.Module):
def __init__(
self, dim: int, elementwise_affine: bool = False, eps: float = 1e-6, device=None, dtype=None
):
"""
Initialize the RMSNorm normalization layer.
Args:
dim (int): The dimension of the input tensor.
eps (float, optional): A small value added to the denominator for numerical stability. Default is 1e-6.
Attributes:
eps (float): A small value added to the denominator for numerical stability.
weight (nn.Parameter): Learnable scaling parameter.
"""
super().__init__()
self.eps = eps
self.learnable_scale = elementwise_affine
if self.learnable_scale:
self.weight = nn.Parameter(torch.empty(dim, device=device, dtype=dtype))
else:
self.register_parameter("weight", None)
def _norm(self, x):
"""
Apply the RMSNorm normalization to the input tensor.
Args:
x (torch.Tensor): The input tensor.
Returns:
torch.Tensor: The normalized tensor.
"""
return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)
def forward(self, x):
"""
Forward pass through the RMSNorm layer.
Args:
x (torch.Tensor): The input tensor.
Returns:
torch.Tensor: The output tensor after applying RMSNorm.
"""
x = self._norm(x)
if self.learnable_scale:
return x * self.weight.to(device=x.device, dtype=x.dtype)
else:
return x
class SwiGLUFeedForward(nn.Module):
def __init__(
self,
dim: int,
hidden_dim: int,
multiple_of: int,
ffn_dim_multiplier: Optional[float] = None,
):
"""
Initialize the FeedForward module.
Args:
dim (int): Input dimension.
hidden_dim (int): Hidden dimension of the feedforward layer.
multiple_of (int): Value to ensure hidden dimension is a multiple of this value.
ffn_dim_multiplier (float, optional): Custom multiplier for hidden dimension. Defaults to None.
Attributes:
w1 (ColumnParallelLinear): Linear transformation for the first layer.
w2 (RowParallelLinear): Linear transformation for the second layer.
w3 (ColumnParallelLinear): Linear transformation for the third layer.
"""
super().__init__()
hidden_dim = int(2 * hidden_dim / 3)
# custom dim factor multiplier
if ffn_dim_multiplier is not None:
hidden_dim = int(ffn_dim_multiplier * hidden_dim)
hidden_dim = multiple_of * ((hidden_dim + multiple_of - 1) // multiple_of)
self.w1 = nn.Linear(dim, hidden_dim, bias=False)
self.w2 = nn.Linear(hidden_dim, dim, bias=False)
self.w3 = nn.Linear(dim, hidden_dim, bias=False)
def forward(self, x):
return self.w2(nn.functional.silu(self.w1(x)) * self.w3(x))
class DismantledBlock(nn.Module):
"""A DiT block with gated adaptive layer norm (adaLN) conditioning."""
ATTENTION_MODES = ("xformers", "torch", "torch-hb", "math", "debug")
def __init__(
self,
hidden_size: int,
num_heads: int,
mlp_ratio: float = 4.0,
attn_mode: str = "xformers",
qkv_bias: bool = False,
pre_only: bool = False,
rmsnorm: bool = False,
scale_mod_only: bool = False,
swiglu: bool = False,
qk_norm: Optional[str] = None,
dtype=None,
device=None,
**block_kwargs,
):
super().__init__()
assert attn_mode in self.ATTENTION_MODES
if not rmsnorm:
self.norm1 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6, dtype=dtype, device=device)
else:
self.norm1 = RMSNorm(hidden_size, elementwise_affine=False, eps=1e-6)
self.attn = SelfAttention(dim=hidden_size, num_heads=num_heads, qkv_bias=qkv_bias, attn_mode=attn_mode, pre_only=pre_only, qk_norm=qk_norm, rmsnorm=rmsnorm, dtype=dtype, device=device)
if not pre_only:
if not rmsnorm:
self.norm2 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6, dtype=dtype, device=device)
else:
self.norm2 = RMSNorm(hidden_size, elementwise_affine=False, eps=1e-6)
mlp_hidden_dim = int(hidden_size * mlp_ratio)
if not pre_only:
if not swiglu:
self.mlp = Mlp(in_features=hidden_size, hidden_features=mlp_hidden_dim, act_layer=nn.GELU(approximate="tanh"), dtype=dtype, device=device)
else:
self.mlp = SwiGLUFeedForward(dim=hidden_size, hidden_dim=mlp_hidden_dim, multiple_of=256)
self.scale_mod_only = scale_mod_only
if not scale_mod_only:
n_mods = 6 if not pre_only else 2
else:
n_mods = 4 if not pre_only else 1
self.adaLN_modulation = nn.Sequential(nn.SiLU(), nn.Linear(hidden_size, n_mods * hidden_size, bias=True, dtype=dtype, device=device))
self.pre_only = pre_only
def pre_attention(self, x: torch.Tensor, c: torch.Tensor):
assert x is not None, "pre_attention called with None input"
if not self.pre_only:
if not self.scale_mod_only:
shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.adaLN_modulation(c).chunk(6, dim=1)
else:
shift_msa = None
shift_mlp = None
scale_msa, gate_msa, scale_mlp, gate_mlp = self.adaLN_modulation(c).chunk(4, dim=1)
qkv = self.attn.pre_attention(modulate(self.norm1(x), shift_msa, scale_msa))
return qkv, (x, gate_msa, shift_mlp, scale_mlp, gate_mlp)
else:
if not self.scale_mod_only:
shift_msa, scale_msa = self.adaLN_modulation(c).chunk(2, dim=1)
else:
shift_msa = None
scale_msa = self.adaLN_modulation(c)
qkv = self.attn.pre_attention(modulate(self.norm1(x), shift_msa, scale_msa))
return qkv, None
def post_attention(self, attn, x, gate_msa, shift_mlp, scale_mlp, gate_mlp):
assert not self.pre_only
x = x + gate_msa.unsqueeze(1) * self.attn.post_attention(attn)
x = x + gate_mlp.unsqueeze(1) * self.mlp(modulate(self.norm2(x), shift_mlp, scale_mlp))
return x
def forward(self, x: torch.Tensor, c: torch.Tensor) -> torch.Tensor:
assert not self.pre_only
(q, k, v), intermediates = self.pre_attention(x, c)
attn = attention(q, k, v, self.attn.num_heads)
return self.post_attention(attn, *intermediates)
def block_mixing(context, x, context_block, x_block, c):
assert context is not None, "block_mixing called with None context"
context_qkv, context_intermediates = context_block.pre_attention(context, c)
x_qkv, x_intermediates = x_block.pre_attention(x, c)
o = []
for t in range(3):
o.append(torch.cat((context_qkv[t], x_qkv[t]), dim=1))
q, k, v = tuple(o)
attn = attention(q, k, v, x_block.attn.num_heads)
context_attn, x_attn = (attn[:, : context_qkv[0].shape[1]], attn[:, context_qkv[0].shape[1] :])
if not context_block.pre_only:
context = context_block.post_attention(context_attn, *context_intermediates)
else:
context = None
x = x_block.post_attention(x_attn, *x_intermediates)
return context, x
class JointBlock(nn.Module):
"""just a small wrapper to serve as a fsdp unit"""
def __init__(self, *args, **kwargs):
super().__init__()
pre_only = kwargs.pop("pre_only")
qk_norm = kwargs.pop("qk_norm", None)
self.context_block = DismantledBlock(*args, pre_only=pre_only, qk_norm=qk_norm, **kwargs)
self.x_block = DismantledBlock(*args, pre_only=False, qk_norm=qk_norm, **kwargs)
def forward(self, *args, **kwargs):
return block_mixing(*args, context_block=self.context_block, x_block=self.x_block, **kwargs)
class FinalLayer(nn.Module):
"""
The final layer of DiT.
"""
def __init__(self, hidden_size: int, patch_size: int, out_channels: int, total_out_channels: Optional[int] = None, dtype=None, device=None):
super().__init__()
self.norm_final = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6, dtype=dtype, device=device)
self.linear = (
nn.Linear(hidden_size, patch_size * patch_size * out_channels, bias=True, dtype=dtype, device=device)
if (total_out_channels is None)
else nn.Linear(hidden_size, total_out_channels, bias=True, dtype=dtype, device=device)
)
self.adaLN_modulation = nn.Sequential(nn.SiLU(), nn.Linear(hidden_size, 2 * hidden_size, bias=True, dtype=dtype, device=device))
def forward(self, x: torch.Tensor, c: torch.Tensor) -> torch.Tensor:
shift, scale = self.adaLN_modulation(c).chunk(2, dim=1)
x = modulate(self.norm_final(x), shift, scale)
x = self.linear(x)
return x
class MMDiT(nn.Module):
"""Diffusion model with a Transformer backbone."""
def __init__(
self,
input_size: int = 32,
patch_size: int = 2,
in_channels: int = 4,
depth: int = 28,
mlp_ratio: float = 4.0,
learn_sigma: bool = False,
adm_in_channels: Optional[int] = None,
context_embedder_config: Optional[Dict] = None,
register_length: int = 0,
attn_mode: str = "torch",
rmsnorm: bool = False,
scale_mod_only: bool = False,
swiglu: bool = False,
out_channels: Optional[int] = None,
pos_embed_scaling_factor: Optional[float] = None,
pos_embed_offset: Optional[float] = None,
pos_embed_max_size: Optional[int] = None,
num_patches = None,
qk_norm: Optional[str] = None,
qkv_bias: bool = True,
dtype = None,
device = None,
):
super().__init__()
self.dtype = dtype
self.learn_sigma = learn_sigma
self.in_channels = in_channels
default_out_channels = in_channels * 2 if learn_sigma else in_channels
self.out_channels = out_channels if out_channels is not None else default_out_channels
self.patch_size = patch_size
self.pos_embed_scaling_factor = pos_embed_scaling_factor
self.pos_embed_offset = pos_embed_offset
self.pos_embed_max_size = pos_embed_max_size
# apply magic --> this defines a head_size of 64
hidden_size = 64 * depth
num_heads = depth
self.num_heads = num_heads
self.x_embedder = PatchEmbed(input_size, patch_size, in_channels, hidden_size, bias=True, strict_img_size=self.pos_embed_max_size is None, dtype=dtype, device=device)
self.t_embedder = TimestepEmbedder(hidden_size, dtype=dtype, device=device)
if adm_in_channels is not None:
assert isinstance(adm_in_channels, int)
self.y_embedder = VectorEmbedder(adm_in_channels, hidden_size, dtype=dtype, device=device)
self.context_embedder = nn.Identity()
if context_embedder_config is not None:
if context_embedder_config["target"] == "torch.nn.Linear":
self.context_embedder = nn.Linear(**context_embedder_config["params"], dtype=dtype, device=device)
self.register_length = register_length
if self.register_length > 0:
self.register = nn.Parameter(torch.randn(1, register_length, hidden_size, dtype=dtype, device=device))
# num_patches = self.x_embedder.num_patches
# Will use fixed sin-cos embedding:
# just use a buffer already
if num_patches is not None:
self.register_buffer(
"pos_embed",
torch.zeros(1, num_patches, hidden_size, dtype=dtype, device=device),
)
else:
self.pos_embed = None
self.joint_blocks = nn.ModuleList(
[
JointBlock(hidden_size, num_heads, mlp_ratio=mlp_ratio, qkv_bias=qkv_bias, attn_mode=attn_mode, pre_only=i == depth - 1, rmsnorm=rmsnorm, scale_mod_only=scale_mod_only, swiglu=swiglu, qk_norm=qk_norm, dtype=dtype, device=device)
for i in range(depth)
]
)
self.final_layer = FinalLayer(hidden_size, patch_size, self.out_channels, dtype=dtype, device=device)
def cropped_pos_embed(self, hw):
assert self.pos_embed_max_size is not None
p = self.x_embedder.patch_size[0]
h, w = hw
# patched size
h = h // p
w = w // p
assert h <= self.pos_embed_max_size, (h, self.pos_embed_max_size)
assert w <= self.pos_embed_max_size, (w, self.pos_embed_max_size)
top = (self.pos_embed_max_size - h) // 2
left = (self.pos_embed_max_size - w) // 2
spatial_pos_embed = rearrange(
self.pos_embed,
"1 (h w) c -> 1 h w c",
h=self.pos_embed_max_size,
w=self.pos_embed_max_size,
)
spatial_pos_embed = spatial_pos_embed[:, top : top + h, left : left + w, :]
spatial_pos_embed = rearrange(spatial_pos_embed, "1 h w c -> 1 (h w) c")
return spatial_pos_embed
def unpatchify(self, x, hw=None):
"""
x: (N, T, patch_size**2 * C)
imgs: (N, H, W, C)
"""
c = self.out_channels
p = self.x_embedder.patch_size[0]
if hw is None:
h = w = int(x.shape[1] ** 0.5)
else:
h, w = hw
h = h // p
w = w // p
assert h * w == x.shape[1]
x = x.reshape(shape=(x.shape[0], h, w, p, p, c))
x = torch.einsum("nhwpqc->nchpwq", x)
imgs = x.reshape(shape=(x.shape[0], c, h * p, w * p))
return imgs
def forward_core_with_concat(self, x: torch.Tensor, c_mod: torch.Tensor, context: Optional[torch.Tensor] = None) -> torch.Tensor:
if self.register_length > 0:
context = torch.cat((repeat(self.register, "1 ... -> b ...", b=x.shape[0]), context if context is not None else torch.Tensor([]).type_as(x)), 1)
# context is B, L', D
# x is B, L, D
for block in self.joint_blocks:
context, x = block(context, x, c=c_mod)
x = self.final_layer(x, c_mod) # (N, T, patch_size ** 2 * out_channels)
return x
def forward(self, x: torch.Tensor, t: torch.Tensor, y: Optional[torch.Tensor] = None, context: Optional[torch.Tensor] = None) -> torch.Tensor:
"""
Forward pass of DiT.
x: (N, C, H, W) tensor of spatial inputs (images or latent representations of images)
t: (N,) tensor of diffusion timesteps
y: (N,) tensor of class labels
"""
hw = x.shape[-2:]
x = self.x_embedder(x) + self.cropped_pos_embed(hw)
c = self.t_embedder(t, dtype=x.dtype) # (N, D)
if y is not None:
y = self.y_embedder(y) # (N, D)
c = c + y # (N, D)
context = self.context_embedder(context)
x = self.forward_core_with_concat(x, c, context)
x = self.unpatchify(x, hw=hw) # (N, out_channels, H, W)
return x
demo/Diffusion/demo_diffusion/utils_sd3/other_impls.py
+555
View File
@@ -0,0 +1,555 @@
# MIT License
# Copyright (c) 2024 Stability AI
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
# The above copyright notice and this permission notice shall be included in all
# copies or substantial portions of the Software.
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.
import torch, math
import numpy as np
from torch import nn
from transformers import CLIPTokenizer, T5TokenizerFast
def load_into(f, model, prefix, device, dtype=None):
"""Just a debugging-friendly hack to apply the weights in a safetensors file to the pytorch module."""
for key in f.keys():
if key.startswith(prefix) and not key.startswith("loss."):
path = key[len(prefix):].split(".")
obj = model
for p in path:
if obj is list:
obj = obj[int(p)]
else:
obj = getattr(obj, p, None)
if obj is None:
print(f"Skipping key '{key}' in safetensors file as '{p}' does not exist in python model")
break
if obj is None:
continue
try:
tensor = f.get_tensor(key).to(device=device)
if dtype is not None:
tensor = tensor.to(dtype=dtype)
obj.requires_grad_(False)
obj.set_(tensor)
except Exception as e:
print(f"Failed to load key '{key}' in safetensors file: {e}")
raise e
def preprocess_image_sd3(image):
image.convert("RGB")
image_np = np.array(image).astype(np.float32) / 255.0
image_np = np.moveaxis(image_np, 2, 0)
batch_images = np.expand_dims(image_np, axis=0).repeat(1, axis=0)
image_torch = torch.from_numpy(batch_images)
image_torch = 2.0 * image_torch - 1.0
return image_torch
#################################################################################################
### Core/Utility
#################################################################################################
def attention(q, k, v, heads, mask=None):
"""Convenience wrapper around a basic attention operation"""
b, _, dim_head = q.shape
dim_head //= heads
q, k, v = map(lambda t: t.view(b, -1, heads, dim_head).transpose(1, 2), (q, k, v))
out = torch.nn.functional.scaled_dot_product_attention(q, k, v, attn_mask=mask, dropout_p=0.0, is_causal=False)
return out.transpose(1, 2).reshape(b, -1, heads * dim_head)
class Mlp(nn.Module):
""" MLP as used in Vision Transformer, MLP-Mixer and related networks"""
def __init__(self, in_features, hidden_features=None, out_features=None, act_layer=nn.GELU, bias=True, dtype=None, device=None):
super().__init__()
out_features = out_features or in_features
hidden_features = hidden_features or in_features
self.fc1 = nn.Linear(in_features, hidden_features, bias=bias, dtype=dtype, device=device)
self.act = act_layer
self.fc2 = nn.Linear(hidden_features, out_features, bias=bias, dtype=dtype, device=device)
def forward(self, x):
x = self.fc1(x)
x = self.act(x)
x = self.fc2(x)
return x
#################################################################################################
### CLIP
#################################################################################################
class CLIPAttention(torch.nn.Module):
def __init__(self, embed_dim, heads, dtype, device):
super().__init__()
self.heads = heads
self.q_proj = nn.Linear(embed_dim, embed_dim, bias=True, dtype=dtype, device=device)
self.k_proj = nn.Linear(embed_dim, embed_dim, bias=True, dtype=dtype, device=device)
self.v_proj = nn.Linear(embed_dim, embed_dim, bias=True, dtype=dtype, device=device)
self.out_proj = nn.Linear(embed_dim, embed_dim, bias=True, dtype=dtype, device=device)
def forward(self, x, mask=None):
q = self.q_proj(x)
k = self.k_proj(x)
v = self.v_proj(x)
out = attention(q, k, v, self.heads, mask)
return self.out_proj(out)
ACTIVATIONS = {
"quick_gelu": lambda a: a * torch.sigmoid(1.702 * a),
"gelu": torch.nn.functional.gelu,
}
class CLIPLayer(torch.nn.Module):
def __init__(self, embed_dim, heads, intermediate_size, intermediate_activation, dtype, device):
super().__init__()
self.layer_norm1 = nn.LayerNorm(embed_dim, dtype=dtype, device=device)
self.self_attn = CLIPAttention(embed_dim, heads, dtype, device)
self.layer_norm2 = nn.LayerNorm(embed_dim, dtype=dtype, device=device)
#self.mlp = CLIPMLP(embed_dim, intermediate_size, intermediate_activation, dtype, device)
self.mlp = Mlp(embed_dim, intermediate_size, embed_dim, act_layer=ACTIVATIONS[intermediate_activation], dtype=dtype, device=device)
def forward(self, x, mask=None):
x += self.self_attn(self.layer_norm1(x), mask)
x += self.mlp(self.layer_norm2(x))
return x
class CLIPEncoder(torch.nn.Module):
def __init__(self, num_layers, embed_dim, heads, intermediate_size, intermediate_activation, dtype, device):
super().__init__()
self.layers = torch.nn.ModuleList([CLIPLayer(embed_dim, heads, intermediate_size, intermediate_activation, dtype, device) for i in range(num_layers)])
def forward(self, x, mask=None, intermediate_output=None):
if intermediate_output is not None:
if intermediate_output < 0:
intermediate_output = len(self.layers) + intermediate_output
intermediate = None
for i, l in enumerate(self.layers):
x = l(x, mask)
if i == intermediate_output:
intermediate = x.clone()
return x, intermediate
class CLIPEmbeddings(torch.nn.Module):
def __init__(self, embed_dim, vocab_size=49408, num_positions=77, dtype=None, device=None):
super().__init__()
self.token_embedding = torch.nn.Embedding(vocab_size, embed_dim, dtype=dtype, device=device)
self.position_embedding = torch.nn.Embedding(num_positions, embed_dim, dtype=dtype, device=device)
def forward(self, input_tokens):
return self.token_embedding(input_tokens) + self.position_embedding.weight
class CLIPTextModel_(torch.nn.Module):
def __init__(self, config_dict, dtype, device):
num_layers = config_dict["num_hidden_layers"]
embed_dim = config_dict["hidden_size"]
heads = config_dict["num_attention_heads"]
intermediate_size = config_dict["intermediate_size"]
intermediate_activation = config_dict["hidden_act"]
super().__init__()
self.embeddings = CLIPEmbeddings(embed_dim, dtype=dtype, device=device)
self.encoder = CLIPEncoder(num_layers, embed_dim, heads, intermediate_size, intermediate_activation, dtype, device)
self.final_layer_norm = nn.LayerNorm(embed_dim, dtype=dtype, device=device)
def forward(self, input_tokens, intermediate_output=None, final_layer_norm_intermediate=True):
x = self.embeddings(input_tokens)
causal_mask = torch.empty(x.shape[1], x.shape[1], dtype=x.dtype, device=x.device).fill_(float("-inf")).triu_(1)
x, i = self.encoder(x, mask=causal_mask, intermediate_output=intermediate_output)
x = self.final_layer_norm(x)
if i is not None and final_layer_norm_intermediate:
i = self.final_layer_norm(i)
pooled_output = x[torch.arange(x.shape[0], device=x.device), input_tokens.to(dtype=torch.int, device=x.device).argmax(dim=-1),]
return x, i, pooled_output
class CLIPTextModel(torch.nn.Module):
def __init__(self, config_dict, dtype, device):
super().__init__()
self.num_layers = config_dict["num_hidden_layers"]
self.text_model = CLIPTextModel_(config_dict, dtype, device)
embed_dim = config_dict["hidden_size"]
self.text_projection = nn.Linear(embed_dim, embed_dim, bias=False, dtype=dtype, device=device)
# WAR for RuntimeError: a leaf Variable that requires grad is being used in an in-place operation.
with torch.no_grad():
self.text_projection.weight.copy_(torch.eye(embed_dim))
self.dtype = dtype
def get_input_embeddings(self):
return self.text_model.embeddings.token_embedding
def set_input_embeddings(self, embeddings):
self.text_model.embeddings.token_embedding = embeddings
def forward(self, *args, **kwargs):
x = self.text_model(*args, **kwargs)
out = self.text_projection(x[2])
return (x[0], x[1], out, x[2])
class SDTokenizer:
def __init__(self, max_length=77, pad_with_end=True, tokenizer=None, has_start_token=True, pad_to_max_length=True, min_length=None):
self.tokenizer = tokenizer
self.max_length = max_length
self.min_length = min_length
empty = self.tokenizer('')["input_ids"]
if has_start_token:
self.tokens_start = 1
self.start_token = empty[0]
self.end_token = empty[1]
else:
self.tokens_start = 0
self.start_token = None
self.end_token = empty[0]
self.pad_with_end = pad_with_end
self.pad_to_max_length = pad_to_max_length
vocab = self.tokenizer.get_vocab()
self.inv_vocab = {v: k for k, v in vocab.items()}
self.max_word_length = 8
def tokenize_with_weights(self, text:str):
"""Tokenize the text, with weight values - presume 1.0 for all and ignore other features here. The details aren't relevant for a reference impl, and weights themselves has weak effect on SD3."""
if self.pad_with_end:
pad_token = self.end_token
else:
pad_token = 0
batch = []
if self.start_token is not None:
batch.append((self.start_token, 1.0))
to_tokenize = text.replace("\n", " ").split(' ')
to_tokenize = [x for x in to_tokenize if x != ""]
for word in to_tokenize:
batch.extend([(t, 1) for t in self.tokenizer(word)["input_ids"][self.tokens_start:-1]])
batch.append((self.end_token, 1.0))
if self.pad_to_max_length:
batch.extend([(pad_token, 1.0)] * (self.max_length - len(batch)))
if self.min_length is not None and len(batch) < self.min_length:
batch.extend([(pad_token, 1.0)] * (self.min_length - len(batch)))
if len(batch) > self.max_length:
batch = batch[:self.max_length]
return [batch]
class SDXLClipGTokenizer(SDTokenizer):
def __init__(self, tokenizer):
super().__init__(pad_with_end=False, tokenizer=tokenizer)
class SD3Tokenizer:
def __init__(self):
clip_tokenizer = CLIPTokenizer.from_pretrained("openai/clip-vit-large-patch14")
self.clip_l = SDTokenizer(tokenizer=clip_tokenizer)
self.clip_g = SDXLClipGTokenizer(clip_tokenizer)
self.t5xxl = T5XXLTokenizer()
def tokenize_with_weights(self, text:str):
out = {}
out["g"] = self.clip_g.tokenize_with_weights(text)
out["l"] = self.clip_l.tokenize_with_weights(text)
out["t5xxl"] = self.t5xxl.tokenize_with_weights(text)
return out
class ClipTokenWeightEncoder:
def encode_token_weights(self, token_weight_pairs):
tokens = list(map(lambda a: a[0], token_weight_pairs[0]))
# model inference
tokens = torch.tensor([tokens], dtype=torch.int64, device="cuda")
out, pooled = self(tokens)
if pooled is not None:
first_pooled = pooled[0:1].cuda()
else:
first_pooled = pooled
output = [out[0:1]]
return torch.cat(output, dim=-2).cuda(), first_pooled
class SDClipModel(torch.nn.Module, ClipTokenWeightEncoder):
"""Uses the CLIP transformer encoder for text (from huggingface)"""
LAYERS = ["last", "pooled", "hidden"]
def __init__(self, device="cuda", max_length=77, layer="last", layer_idx=None, textmodel_json_config=None, dtype=None, model_class=CLIPTextModel,
special_tokens={"start": 49406, "end": 49407, "pad": 49407}, layer_norm_hidden_state=True, return_projected_pooled=True):
super().__init__()
assert layer in self.LAYERS
self.transformer = model_class(textmodel_json_config, dtype, device)
self.num_layers = self.transformer.num_layers
self.max_length = max_length
self.transformer = self.transformer.eval()
for param in self.parameters():
param.requires_grad = False
self.layer = layer
self.layer_idx = None
self.special_tokens = special_tokens
self.logit_scale = torch.nn.Parameter(torch.tensor(4.6055))
self.layer_norm_hidden_state = layer_norm_hidden_state
self.return_projected_pooled = return_projected_pooled
if layer == "hidden":
assert layer_idx is not None
assert abs(layer_idx) < self.num_layers
self.set_clip_options({"layer": layer_idx})
self.options_default = (self.layer, self.layer_idx, self.return_projected_pooled)
def set_clip_options(self, options):
layer_idx = options.get("layer", self.layer_idx)
self.return_projected_pooled = options.get("projected_pooled", self.return_projected_pooled)
if layer_idx is None or abs(layer_idx) > self.num_layers:
self.layer = "last"
else:
self.layer = "hidden"
self.layer_idx = layer_idx
def forward(self, tokens):
backup_embeds = self.transformer.get_input_embeddings()
outputs = self.transformer(tokens, intermediate_output=self.layer_idx, final_layer_norm_intermediate=self.layer_norm_hidden_state)
self.transformer.set_input_embeddings(backup_embeds)
if self.layer == "last":
z = outputs[0]
else:
z = outputs[1]
pooled_output = None
if len(outputs) >= 3:
if not self.return_projected_pooled and len(outputs) >= 4 and outputs[3] is not None:
pooled_output = outputs[3].float()
elif outputs[2] is not None:
pooled_output = outputs[2].float()
return z.float(), pooled_output
class SDXLClipG(SDClipModel):
"""Wraps the CLIP-G model into the SD-CLIP-Model interface"""
def __init__(self, config, device="cuda", layer="penultimate", layer_idx=None, dtype=None):
if layer == "penultimate":
layer="hidden"
layer_idx=-2
super().__init__(device=device, layer=layer, layer_idx=layer_idx, textmodel_json_config=config, dtype=dtype, special_tokens={"start": 49406, "end": 49407, "pad": 0}, layer_norm_hidden_state=False)
class T5XXLModel(SDClipModel):
"""Wraps the T5-XXL model into the SD-CLIP-Model interface for convenience"""
def __init__(self, config, device="cuda", layer="last", layer_idx=None, dtype=None):
super().__init__(device=device, layer=layer, layer_idx=layer_idx, textmodel_json_config=config, dtype=dtype, special_tokens={"end": 1, "pad": 0}, model_class=T5)
#################################################################################################
### T5 implementation, for the T5-XXL text encoder portion, largely pulled from upstream impl
#################################################################################################
class T5XXLTokenizer(SDTokenizer):
"""Wraps the T5 Tokenizer from HF into the SDTokenizer interface"""
def __init__(self):
super().__init__(pad_with_end=False, tokenizer=T5TokenizerFast.from_pretrained("google/t5-v1_1-xxl"), has_start_token=False, pad_to_max_length=False, max_length=99999999, min_length=77)
class T5LayerNorm(torch.nn.Module):
def __init__(self, hidden_size, eps=1e-6, dtype=None, device=None):
super().__init__()
self.weight = torch.nn.Parameter(torch.ones(hidden_size, dtype=dtype, device=device))
self.variance_epsilon = eps
def forward(self, x):
variance = x.pow(2).mean(-1, keepdim=True)
x = x * torch.rsqrt(variance + self.variance_epsilon)
return self.weight.to(device=x.device, dtype=x.dtype) * x
class T5DenseGatedActDense(torch.nn.Module):
def __init__(self, model_dim, ff_dim, dtype, device):
super().__init__()
self.wi_0 = nn.Linear(model_dim, ff_dim, bias=False, dtype=dtype, device=device)
self.wi_1 = nn.Linear(model_dim, ff_dim, bias=False, dtype=dtype, device=device)
self.wo = nn.Linear(ff_dim, model_dim, bias=False, dtype=dtype, device=device)
def forward(self, x):
hidden_gelu = torch.nn.functional.gelu(self.wi_0(x), approximate="tanh")
hidden_linear = self.wi_1(x)
x = hidden_gelu * hidden_linear
x = self.wo(x)
return x
class T5LayerFF(torch.nn.Module):
def __init__(self, model_dim, ff_dim, dtype, device):
super().__init__()
self.DenseReluDense = T5DenseGatedActDense(model_dim, ff_dim, dtype, device)
self.layer_norm = T5LayerNorm(model_dim, dtype=dtype, device=device)
def forward(self, x):
forwarded_states = self.layer_norm(x)
forwarded_states = self.DenseReluDense(forwarded_states)
x += forwarded_states
return x
class T5Attention(torch.nn.Module):
def __init__(self, model_dim, inner_dim, num_heads, relative_attention_bias, dtype, device):
super().__init__()
# Mesh TensorFlow initialization to avoid scaling before softmax
self.q = nn.Linear(model_dim, inner_dim, bias=False, dtype=dtype, device=device)
self.k = nn.Linear(model_dim, inner_dim, bias=False, dtype=dtype, device=device)
self.v = nn.Linear(model_dim, inner_dim, bias=False, dtype=dtype, device=device)
self.o = nn.Linear(inner_dim, model_dim, bias=False, dtype=dtype, device=device)
self.num_heads = num_heads
self.relative_attention_bias = None
if relative_attention_bias:
self.relative_attention_num_buckets = 32
self.relative_attention_max_distance = 128
self.relative_attention_bias = torch.nn.Embedding(self.relative_attention_num_buckets, self.num_heads, device=device, dtype=dtype)
@staticmethod
def _relative_position_bucket(relative_position, bidirectional=True, num_buckets=32, max_distance=128):
"""
Adapted from Mesh Tensorflow:
https://github.com/tensorflow/mesh/blob/0cb87fe07da627bf0b7e60475d59f95ed6b5be3d/mesh_tensorflow/transformer/transformer_layers.py#L593
Translate relative position to a bucket number for relative attention. The relative position is defined as
memory_position - query_position, i.e. the distance in tokens from the attending position to the attended-to
position. If bidirectional=False, then positive relative positions are invalid. We use smaller buckets for
small absolute relative_position and larger buckets for larger absolute relative_positions. All relative
positions >=max_distance map to the same bucket. All relative positions <=-max_distance map to the same bucket.
This should allow for more graceful generalization to longer sequences than the model has been trained on
Args:
relative_position: an int32 Tensor
bidirectional: a boolean - whether the attention is bidirectional
num_buckets: an integer
max_distance: an integer
Returns:
a Tensor with the same shape as relative_position, containing int32 values in the range [0, num_buckets)
"""
relative_buckets = 0
if bidirectional:
num_buckets //= 2
relative_buckets += (relative_position > 0).to(torch.long) * num_buckets
relative_position = torch.abs(relative_position)
else:
relative_position = -torch.min(relative_position, torch.zeros_like(relative_position))
# now relative_position is in the range [0, inf)
# half of the buckets are for exact increments in positions
max_exact = num_buckets // 2
is_small = relative_position < max_exact
# The other half of the buckets are for logarithmically bigger bins in positions up to max_distance
relative_position_if_large = max_exact + (
torch.log(relative_position.float() / max_exact)
/ math.log(max_distance / max_exact)
* (num_buckets - max_exact)
).to(torch.long)
relative_position_if_large = torch.min(relative_position_if_large, torch.full_like(relative_position_if_large, num_buckets - 1))
relative_buckets += torch.where(is_small, relative_position, relative_position_if_large)
return relative_buckets
def compute_bias(self, query_length, key_length, device):
"""Compute binned relative position bias"""
context_position = torch.arange(query_length, dtype=torch.long, device=device)[:, None]
memory_position = torch.arange(key_length, dtype=torch.long, device=device)[None, :]
relative_position = memory_position - context_position # shape (query_length, key_length)
relative_position_bucket = self._relative_position_bucket(
relative_position, # shape (query_length, key_length)
bidirectional=True,
num_buckets=self.relative_attention_num_buckets,
max_distance=self.relative_attention_max_distance,
)
values = self.relative_attention_bias(relative_position_bucket) # shape (query_length, key_length, num_heads)
values = values.permute([2, 0, 1]).unsqueeze(0) # shape (1, num_heads, query_length, key_length)
return values
def forward(self, x, past_bias=None):
q = self.q(x)
k = self.k(x)
v = self.v(x)
if self.relative_attention_bias is not None:
past_bias = self.compute_bias(x.shape[1], x.shape[1], x.device)
if past_bias is not None:
mask = past_bias
out = attention(q, k * ((k.shape[-1] / self.num_heads) ** 0.5), v, self.num_heads, mask)
return self.o(out), past_bias
class T5LayerSelfAttention(torch.nn.Module):
def __init__(self, model_dim, inner_dim, ff_dim, num_heads, relative_attention_bias, dtype, device):
super().__init__()
self.SelfAttention = T5Attention(model_dim, inner_dim, num_heads, relative_attention_bias, dtype, device)
self.layer_norm = T5LayerNorm(model_dim, dtype=dtype, device=device)
def forward(self, x, past_bias=None):
output, past_bias = self.SelfAttention(self.layer_norm(x), past_bias=past_bias)
x += output
return x, past_bias
class T5Block(torch.nn.Module):
def __init__(self, model_dim, inner_dim, ff_dim, num_heads, relative_attention_bias, dtype, device):
super().__init__()
self.layer = torch.nn.ModuleList()
self.layer.append(T5LayerSelfAttention(model_dim, inner_dim, ff_dim, num_heads, relative_attention_bias, dtype, device))
self.layer.append(T5LayerFF(model_dim, ff_dim, dtype, device))
def forward(self, x, past_bias=None):
x, past_bias = self.layer[0](x, past_bias)
x = self.layer[-1](x)
return x, past_bias
class T5Stack(torch.nn.Module):
def __init__(self, num_layers, model_dim, inner_dim, ff_dim, num_heads, vocab_size, dtype, device):
super().__init__()
self.embed_tokens = torch.nn.Embedding(vocab_size, model_dim, device=device, dtype=dtype)
self.block = torch.nn.ModuleList([T5Block(model_dim, inner_dim, ff_dim, num_heads, relative_attention_bias=(i == 0), dtype=dtype, device=device) for i in range(num_layers)])
self.final_layer_norm = T5LayerNorm(model_dim, dtype=dtype, device=device)
def forward(self, input_ids, intermediate_output=None, final_layer_norm_intermediate=True):
intermediate = None
x = self.embed_tokens(input_ids)
past_bias = None
for i, l in enumerate(self.block):
x, past_bias = l(x, past_bias)
if i == intermediate_output:
intermediate = x.clone()
x = self.final_layer_norm(x)
if intermediate is not None and final_layer_norm_intermediate:
intermediate = self.final_layer_norm(intermediate)
return x, intermediate
class T5(torch.nn.Module):
def __init__(self, config_dict, dtype, device):
super().__init__()
self.num_layers = config_dict["num_layers"]
self.encoder = T5Stack(self.num_layers, config_dict["d_model"], config_dict["d_model"], config_dict["d_ff"], config_dict["num_heads"], config_dict["vocab_size"], dtype, device)
self.dtype = dtype
def get_input_embeddings(self):
return self.encoder.embed_tokens
def set_input_embeddings(self, embeddings):
self.encoder.embed_tokens = embeddings
def forward(self, *args, **kwargs):
return self.encoder(*args, **kwargs)
demo/Diffusion/demo_diffusion/utils_sd3/sd3_impls.py
+391
View File
@@ -0,0 +1,391 @@
# MIT License
# Copyright (c) 2024 Stability AI
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
# The above copyright notice and this permission notice shall be included in all
# copies or substantial portions of the Software.
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.
import math
import einops
import torch
from PIL import Image
from demo_diffusion.utils_sd3.mmdit import MMDiT
#################################################################################################
### MMDiT Model Wrapping
#################################################################################################
class ModelSamplingDiscreteFlow(torch.nn.Module):
"""Helper for sampler scheduling (ie timestep/sigma calculations) for Discrete Flow models"""
def __init__(self, shift=1.0):
super().__init__()
self.shift = shift
timesteps = 1000
ts = self.sigma(torch.arange(1, timesteps + 1, 1))
self.register_buffer('sigmas', ts)
@property
def sigma_min(self):
return self.sigmas[0]
@property
def sigma_max(self):
return self.sigmas[-1]
def timestep(self, sigma):
return sigma * 1000
def sigma(self, timestep: torch.Tensor):
timestep = timestep / 1000.0
if self.shift == 1.0:
return timestep
return self.shift * timestep / (1 + (self.shift - 1) * timestep)
def calculate_denoised(self, sigma, model_output, model_input):
sigma = sigma.view(sigma.shape[:1] + (1,) * (model_output.ndim - 1))
return model_input - model_output * sigma
def noise_scaling(self, sigma, noise, latent_image, max_denoise=False):
return sigma * noise + (1.0 - sigma) * latent_image
class BaseModel(torch.nn.Module):
"""Wrapper around the core MM-DiT model"""
def __init__(self, shift=1.0, device=None, dtype=torch.float32, file=None, prefix=""):
super().__init__()
# Important configuration values can be quickly determined by checking shapes in the source file
# Some of these will vary between models (eg 2B vs 8B primarily differ in their depth, but also other details change)
patch_size = file.get_tensor(f"{prefix}x_embedder.proj.weight").shape[2]
depth = file.get_tensor(f"{prefix}x_embedder.proj.weight").shape[0] // 64
num_patches = file.get_tensor(f"{prefix}pos_embed").shape[1]
pos_embed_max_size = round(math.sqrt(num_patches))
adm_in_channels = file.get_tensor(f"{prefix}y_embedder.mlp.0.weight").shape[1]
context_shape = file.get_tensor(f"{prefix}context_embedder.weight").shape
context_embedder_config = {
"target": "torch.nn.Linear",
"params": {
"in_features": context_shape[1],
"out_features": context_shape[0]
}
}
self.diffusion_model = MMDiT(input_size=None, pos_embed_scaling_factor=None, pos_embed_offset=None, pos_embed_max_size=pos_embed_max_size, patch_size=patch_size, in_channels=16, depth=depth, num_patches=num_patches, adm_in_channels=adm_in_channels, context_embedder_config=context_embedder_config, device=device, dtype=dtype)
self.model_sampling = ModelSamplingDiscreteFlow(shift=shift)
def forward(self, x, sigma, c_crossattn=None, y=None):
dtype = self.get_dtype()
timestep = self.model_sampling.timestep(sigma).float()
model_output = self.diffusion_model(x.to(dtype), timestep, context=c_crossattn.to(dtype), y=y.to(dtype)).float()
return self.model_sampling.calculate_denoised(sigma, model_output, x)
def get_dtype(self):
return self.diffusion_model.dtype
class CFGDenoiser(torch.nn.Module):
"""Helper for applying CFG Scaling to diffusion outputs"""
def __init__(self, model):
super().__init__()
self.model = model
def forward(self, x, timestep, cond, uncond, cond_scale):
# Run cond and uncond in a batch together
batched = self.model(torch.cat([x, x]), torch.cat([timestep, timestep]), c_crossattn=torch.cat([cond["c_crossattn"], uncond["c_crossattn"]]), y=torch.cat([cond["y"], uncond["y"]]))
# Then split and apply CFG Scaling
pos_out, neg_out = batched.chunk(2)
scaled = neg_out + (pos_out - neg_out) * cond_scale
return scaled
class SD3LatentFormat:
"""Latents are slightly shifted from center - this class must be called after VAE Decode to correct for the shift"""
def __init__(self):
self.scale_factor = 1.5305
self.shift_factor = 0.0609
def process_in(self, latent):
return (latent - self.shift_factor) * self.scale_factor
def process_out(self, latent):
return (latent / self.scale_factor) + self.shift_factor
def decode_latent_to_preview(self, x0):
"""Quick RGB approximate preview of sd3 latents"""
factors = torch.tensor([
[-0.0645, 0.0177, 0.1052], [ 0.0028, 0.0312, 0.0650],
[ 0.1848, 0.0762, 0.0360], [ 0.0944, 0.0360, 0.0889],
[ 0.0897, 0.0506, -0.0364], [-0.0020, 0.1203, 0.0284],
[ 0.0855, 0.0118, 0.0283], [-0.0539, 0.0658, 0.1047],
[-0.0057, 0.0116, 0.0700], [-0.0412, 0.0281, -0.0039],
[ 0.1106, 0.1171, 0.1220], [-0.0248, 0.0682, -0.0481],
[ 0.0815, 0.0846, 0.1207], [-0.0120, -0.0055, -0.0867],
[-0.0749, -0.0634, -0.0456], [-0.1418, -0.1457, -0.1259]
], device="cuda")
latent_image = x0[0].permute(1, 2, 0).cuda() @ factors
latents_ubyte = (((latent_image + 1) / 2)
.clamp(0, 1) # change scale from -1..1 to 0..1
.mul(0xFF) # to 0..255
.byte()).cuda()
return Image.fromarray(latents_ubyte.numpy())
#################################################################################################
### K-Diffusion Sampling
#################################################################################################
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
return x[(...,) + (None,) * dims_to_append]
def to_d(x, sigma, denoised):
"""Converts a denoiser output to a Karras ODE derivative."""
return (x - denoised) / append_dims(sigma, x.ndim)
@torch.no_grad()
@torch.autocast("cuda", dtype=torch.float16)
def sample_euler(func, x, sigmas, extra_args=None):
"""Implements Algorithm 2 (Euler steps) from Karras et al. (2022)."""
extra_args = {} if extra_args is None else extra_args
s_in = x.new_ones([x.shape[0]])
for i in range(len(sigmas) - 1):
sigma_hat = sigmas[i]
denoised = func(x, sigma_hat * s_in, **extra_args)
d = to_d(x, sigma_hat, denoised)
dt = sigmas[i + 1] - sigma_hat
# Euler method
x = x + d * dt
return x
#################################################################################################
### VAE
#################################################################################################
def Normalize(in_channels, num_groups=32, dtype=torch.float32, device=None):
return torch.nn.GroupNorm(num_groups=num_groups, num_channels=in_channels, eps=1e-6, affine=True, dtype=dtype, device=device)
class ResnetBlock(torch.nn.Module):
def __init__(self, *, in_channels, out_channels=None, dtype=torch.float32, device=None):
super().__init__()
self.in_channels = in_channels
out_channels = in_channels if out_channels is None else out_channels
self.out_channels = out_channels
self.norm1 = Normalize(in_channels, dtype=dtype, device=device)
self.conv1 = torch.nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=1, padding=1, dtype=dtype, device=device)
self.norm2 = Normalize(out_channels, dtype=dtype, device=device)
self.conv2 = torch.nn.Conv2d(out_channels, out_channels, kernel_size=3, stride=1, padding=1, dtype=dtype, device=device)
if self.in_channels != self.out_channels:
self.nin_shortcut = torch.nn.Conv2d(in_channels, out_channels, kernel_size=1, stride=1, padding=0, dtype=dtype, device=device)
else:
self.nin_shortcut = None
self.swish = torch.nn.SiLU(inplace=True)
def forward(self, x):
hidden = x
hidden = self.norm1(hidden)
hidden = self.swish(hidden)
hidden = self.conv1(hidden)
hidden = self.norm2(hidden)
hidden = self.swish(hidden)
hidden = self.conv2(hidden)
if self.in_channels != self.out_channels:
x = self.nin_shortcut(x)
return x + hidden
class AttnBlock(torch.nn.Module):
def __init__(self, in_channels, dtype=torch.float32, device=None):
super().__init__()
self.norm = Normalize(in_channels, dtype=dtype, device=device)
self.q = torch.nn.Conv2d(in_channels, in_channels, kernel_size=1, stride=1, padding=0, dtype=dtype, device=device)
self.k = torch.nn.Conv2d(in_channels, in_channels, kernel_size=1, stride=1, padding=0, dtype=dtype, device=device)
self.v = torch.nn.Conv2d(in_channels, in_channels, kernel_size=1, stride=1, padding=0, dtype=dtype, device=device)
self.proj_out = torch.nn.Conv2d(in_channels, in_channels, kernel_size=1, stride=1, padding=0, dtype=dtype, device=device)
def forward(self, x):
hidden = self.norm(x)
q = self.q(hidden)
k = self.k(hidden)
v = self.v(hidden)
b, c, h, w = q.shape
q, k, v = map(lambda x: einops.rearrange(x, "b c h w -> b 1 (h w) c").contiguous(), (q, k, v))
hidden = torch.nn.functional.scaled_dot_product_attention(q, k, v) # scale is dim ** -0.5 per default
hidden = einops.rearrange(hidden, "b 1 (h w) c -> b c h w", h=h, w=w, c=c, b=b)
hidden = self.proj_out(hidden)
return x + hidden
class Downsample(torch.nn.Module):
def __init__(self, in_channels, dtype=torch.float32, device=None):
super().__init__()
self.conv = torch.nn.Conv2d(in_channels, in_channels, kernel_size=3, stride=2, padding=0, dtype=dtype, device=device)
def forward(self, x):
pad = (0,1,0,1)
x = torch.nn.functional.pad(x, pad, mode="constant", value=0)
x = self.conv(x)
return x
class Upsample(torch.nn.Module):
def __init__(self, in_channels, dtype=torch.float32, device=None):
super().__init__()
self.conv = torch.nn.Conv2d(in_channels, in_channels, kernel_size=3, stride=1, padding=1, dtype=dtype, device=device)
def forward(self, x):
x = torch.nn.functional.interpolate(x, scale_factor=2.0, mode="nearest")
x = self.conv(x)
return x
class VAEEncoder(torch.nn.Module):
def __init__(self, ch=128, ch_mult=(1,2,4,4), num_res_blocks=2, in_channels=3, z_channels=16, dtype=torch.float32, device=None):
super().__init__()
self.num_resolutions = len(ch_mult)
self.num_res_blocks = num_res_blocks
# downsampling
self.conv_in = torch.nn.Conv2d(in_channels, ch, kernel_size=3, stride=1, padding=1, dtype=dtype, device=device)
in_ch_mult = (1,) + tuple(ch_mult)
self.in_ch_mult = in_ch_mult
self.down = torch.nn.ModuleList()
for i_level in range(self.num_resolutions):
block = torch.nn.ModuleList()
attn = torch.nn.ModuleList()
block_in = ch*in_ch_mult[i_level]
block_out = ch*ch_mult[i_level]
for i_block in range(num_res_blocks):
block.append(ResnetBlock(in_channels=block_in, out_channels=block_out, dtype=dtype, device=device))
block_in = block_out
down = torch.nn.Module()
down.block = block
down.attn = attn
if i_level != self.num_resolutions - 1:
down.downsample = Downsample(block_in, dtype=dtype, device=device)
self.down.append(down)
# middle
self.mid = torch.nn.Module()
self.mid.block_1 = ResnetBlock(in_channels=block_in, out_channels=block_in, dtype=dtype, device=device)
self.mid.attn_1 = AttnBlock(block_in, dtype=dtype, device=device)
self.mid.block_2 = ResnetBlock(in_channels=block_in, out_channels=block_in, dtype=dtype, device=device)
# end
self.norm_out = Normalize(block_in, dtype=dtype, device=device)
self.conv_out = torch.nn.Conv2d(block_in, 2 * z_channels, kernel_size=3, stride=1, padding=1, dtype=dtype, device=device)
self.swish = torch.nn.SiLU(inplace=True)
def forward(self, x):
# downsampling
hs = [self.conv_in(x)]
for i_level in range(self.num_resolutions):
for i_block in range(self.num_res_blocks):
h = self.down[i_level].block[i_block](hs[-1])
hs.append(h)
if i_level != self.num_resolutions-1:
hs.append(self.down[i_level].downsample(hs[-1]))
# middle
h = hs[-1]
h = self.mid.block_1(h)
h = self.mid.attn_1(h)
h = self.mid.block_2(h)
# end
h = self.norm_out(h)
h = self.swish(h)
h = self.conv_out(h)
return h
class VAEDecoder(torch.nn.Module):
def __init__(self, ch=128, out_ch=3, ch_mult=(1, 2, 4, 4), num_res_blocks=2, resolution=256, z_channels=16, dtype=torch.float32, device=None):
super().__init__()
self.num_resolutions = len(ch_mult)
self.num_res_blocks = num_res_blocks
block_in = ch * ch_mult[self.num_resolutions - 1]
curr_res = resolution // 2 ** (self.num_resolutions - 1)
# z to block_in
self.conv_in = torch.nn.Conv2d(z_channels, block_in, kernel_size=3, stride=1, padding=1, dtype=dtype, device=device)
# middle
self.mid = torch.nn.Module()
self.mid.block_1 = ResnetBlock(in_channels=block_in, out_channels=block_in, dtype=dtype, device=device)
self.mid.attn_1 = AttnBlock(block_in, dtype=dtype, device=device)
self.mid.block_2 = ResnetBlock(in_channels=block_in, out_channels=block_in, dtype=dtype, device=device)
# upsampling
self.up = torch.nn.ModuleList()
for i_level in reversed(range(self.num_resolutions)):
block = torch.nn.ModuleList()
block_out = ch * ch_mult[i_level]
for i_block in range(self.num_res_blocks + 1):
block.append(ResnetBlock(in_channels=block_in, out_channels=block_out, dtype=dtype, device=device))
block_in = block_out
up = torch.nn.Module()
up.block = block
if i_level != 0:
up.upsample = Upsample(block_in, dtype=dtype, device=device)
curr_res = curr_res * 2
self.up.insert(0, up) # prepend to get consistent order
# end
self.norm_out = Normalize(block_in, dtype=dtype, device=device)
self.conv_out = torch.nn.Conv2d(block_in, out_ch, kernel_size=3, stride=1, padding=1, dtype=dtype, device=device)
self.swish = torch.nn.SiLU(inplace=True)
def forward(self, z):
# z to block_in
hidden = self.conv_in(z)
# middle
hidden = self.mid.block_1(hidden)
hidden = self.mid.attn_1(hidden)
hidden = self.mid.block_2(hidden)
# upsampling
for i_level in reversed(range(self.num_resolutions)):
for i_block in range(self.num_res_blocks + 1):
hidden = self.up[i_level].block[i_block](hidden)
if i_level != 0:
hidden = self.up[i_level].upsample(hidden)
# end
hidden = self.norm_out(hidden)
hidden = self.swish(hidden)
hidden = self.conv_out(hidden)
return hidden
class SDVAE(torch.nn.Module):
def __init__(self, dtype=torch.float32, device=None):
super().__init__()
self.encoder = VAEEncoder(dtype=dtype, device=device)
self.decoder = VAEDecoder(dtype=dtype, device=device)
@torch.autocast("cuda", dtype=torch.float16)
def decode(self, latent):
return self.decoder(latent)
@torch.autocast("cuda", dtype=torch.float16)
def encode(self, image):
hidden = self.encoder(image)
mean, logvar = torch.chunk(hidden, 2, dim=1)
logvar = torch.clamp(logvar, -30.0, 20.0)
std = torch.exp(0.5 * logvar)
return mean + std * torch.randn_like(mean)
demo/Diffusion/demo_img2img.py
+81
View File
@@ -0,0 +1,81 @@
#
# 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.
#
# Configure dependencies before any external imports
from demo_diffusion import deps
deps.configure("sd")
import argparse
import PIL
from cuda.bindings import runtime as cudart
from PIL import 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 Img2Img Demo")
parser = dd_argparse.add_arguments(parser)
parser.add_argument('--input-image', type=str, default="", help="Path to the input image")
return parser.parse_args()
if __name__ == "__main__":
print("[I] Initializing StableDiffusion img2img demo using TensorRT")
args = parseArgs()
if args.input_image:
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 = image_module.download_image(url)
image_width, image_height = input_image.size
if image_height != args.height or image_width != args.width:
print(f"[I] Resizing input_image to {args.height}x{args.width}")
input_image = input_image.resize((args.width, args.height))
image_height, image_width = args.height, args.width
if isinstance(input_image, PIL.Image.Image):
input_image = image_module.preprocess_image(input_image)
kwargs_init_pipeline, kwargs_load_engine, args_run_demo = dd_argparse.process_pipeline_args(args)
# Initialize demo
demo = pipeline_module.StableDiffusionPipeline(
pipeline_type=pipeline_module.PIPELINE_TYPE.IMG2IMG, **kwargs_init_pipeline
)
# Load TensorRT engines and pytorch modules
demo.loadEngines(
args.engine_dir,
args.framework_model_dir,
args.onnx_dir,
**kwargs_load_engine)
# Load resources
_, shared_device_memory = cudart.cudaMalloc(demo.calculateMaxDeviceMemory())
demo.activateEngines(shared_device_memory)
demo.loadResources(args.height, args.width, args.batch_size, args.seed)
# Run inference
demo_kwargs = {'input_image': input_image, 'image_strength': 0.75}
demo.run(*args_run_demo, **demo_kwargs)
demo.teardown()
demo/Diffusion/demo_img2img_flux.py
Executable
+284
View File
@@ -0,0 +1,284 @@
#
# 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.
#
# Configure dependencies before any external imports
from demo_diffusion import deps
deps.configure("flux")
import argparse
import os
import controlnet_aux
from cuda.bindings import runtime as cudart
from PIL import Image
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 = dd_argparse.add_arguments(parser)
parser.add_argument(
"--version",
type=str,
default="flux.1-dev",
choices=("flux.1-dev", "flux.1-schnell", "flux.1-dev-canny", "flux.1-dev-depth", "flux.1-kontext-dev"),
help="Version of Flux",
)
parser.add_argument(
"--prompt2",
default=None,
nargs="*",
help="Text prompt(s) to be sent to the T5 tokenizer and text encoder. If not defined, prompt will be used instead",
)
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)",
)
parser.add_argument("--denoising-steps", type=int, default=50, help="Number of denoising steps")
parser.add_argument(
"--guidance-scale",
type=float,
default=3.5,
help="Value of classifier-free guidance scale (must be greater than 1)",
)
parser.add_argument(
"--max_sequence_length",
type=int,
help="Maximum sequence length to use with the prompt. Can be up to 512 for the dev and 256 for the schnell variant.",
)
parser.add_argument(
"--t5-ws-percentage",
type=int,
default=None,
help="Set runtime weight streaming budget as the percentage of the size of streamable weights for the T5 model. This argument only takes effect when --ws is set. 0 streams the most weights and 100 or None streams no weights. ",
)
parser.add_argument(
"--transformer-ws-percentage",
type=int,
default=None,
help="Set runtime weight streaming budget as the percentage of the size of streamable weights for the transformer model. This argument only takes effect when --ws is set. 0 streams the most weights and 100 or None streams no weights.",
)
parser.add_argument(
"--control-image",
type=str,
default=None,
help="Path to the control image for the flux.1-dev-canny and flux.1-dev-depth pipelines",
)
parser.add_argument(
"--input-image",
type=str,
default=None,
help="Path to the input conditioning image for the flux.1-dev and flux.1-schnell img2img pipelines",
)
parser.add_argument(
"--kontext-image",
type=str,
default=None,
help="Path to the input image for Kontext pipeline (flux.1-kontext-dev only, required)",
)
parser.add_argument(
"--image-strength",
type=float,
default=1.0,
help="Indicates extent to transform the reference `image`. Must be between 0 and 1. A value of 1 essentially ignores the input image.",
)
parser.add_argument(
"--calibration-dataset",
type=str,
default=None,
help="Path to the calibration dataset for quantization (only enabled for controlnet)",
)
return parser.parse_args()
def process_demo_args(args):
batch_size = args.batch_size
prompt = args.prompt
# If prompt2 is not defined, use prompt instead
prompt2 = args.prompt2 or prompt
# Process input args
if not isinstance(prompt, list):
raise ValueError(f"`prompt` must be of type `list[str]`, but is {type(prompt)}")
prompt = prompt * batch_size
if not isinstance(prompt2, list):
raise ValueError(f"`prompt2` must be of type `str` list, but is {type(prompt2)}")
if len(prompt2) == 1:
prompt2 = prompt2 * batch_size
max_seq_supported_by_model = {
"flux.1-schnell": 256,
"flux.1-dev": 512,
"flux.1-dev-canny": 512,
"flux.1-dev-depth": 512,
"flux.1-kontext-dev": 512,
}[args.version]
if args.max_sequence_length is not None:
if args.max_sequence_length > max_seq_supported_by_model:
raise ValueError(
f"For {args.version}, `max_sequence_length` cannot be greater than {max_seq_supported_by_model} but is {args.max_sequence_length}"
)
else:
args.max_sequence_length = max_seq_supported_by_model
controlnet_type = "depth" if "depth" in args.version else "canny" if "canny" in args.version else ""
if controlnet_type:
if args.input_image:
raise ValueError(
f"--input-image is a valid input for versions [flux.1-dev, flux.1-schnell]. Provided {args.version}"
)
if not args.control_image:
raise ValueError(
"--control-image input is required for versions [flux.1-dev-canny, flux.1-dev-depth]. Please provide it using --control-image flag."
)
args.control_image = Image.open(args.control_image).convert("RGB")
if controlnet_type == "canny":
processor = controlnet_aux.CannyDetector()
args.control_image = processor(
args.control_image, low_threshold=50, high_threshold=200, detect_resolution=1024, image_resolution=1024
)
elif controlnet_type == "depth":
args.control_image = controlnet_aux.LeresDetector.from_pretrained("lllyasviel/Annotators")(
args.control_image
)
else:
raise ValueError("Invalid controlnet type")
else:
if args.control_image:
raise ValueError(
f"--control-image is a valid input for versions [flux.1-dev-canny, flux.1-dev-depth]. Provided {args.version}"
)
# Handle input image for img2img pipelines
if args.version == "flux.1-kontext-dev":
# For Kontext pipeline, only use kontext-image
if not args.kontext_image:
raise ValueError(
"--kontext-image is required for the Kontext pipeline. Please provide it using the --kontext-image flag."
)
if args.input_image:
raise ValueError(
"--input-image is not supported for the Kontext pipeline. Please use --kontext-image instead."
)
# Kontext pipeline doesn't resize the input image
args.kontext_image = Image.open(args.kontext_image).convert("RGB")
else:
if not args.input_image:
raise ValueError(
"--input-image is required for the img2img pipeline. Please provide it using the --input-image flag."
)
args.input_image = Image.open(args.input_image).convert("RGB").resize((args.width, args.height))
if args.fp8:
if args.version == "flux.1-dev" or args.version == "flux.1-schnell":
raise ValueError("--fp8 is currently not supported for Flux.1-dev and Flux.1-schnell img2img pipelines.")
if not args.calibration_dataset:
args.calibration_dataset = os.path.join(f"{controlnet_type}-eval", "benchmark")
print(f"[W] Calibration dataset path not provided, setting default path to {args.calibration_dataset}.")
if not os.path.exists(args.calibration_dataset):
print(
f"[W] Could not find the calibration dataset at {args.calibration_dataset}, and will fallback to using pre-exported ONNX models. Please follow the instructions in README to download calibration dataset and provide the path if pre-exported ONNX models are not provided either."
)
if args.version == "flux.1-kontext-dev" and not args.download_onnx_models:
raise ValueError(
"--download-onnx-models is required when using --fp8 for Flux.1-kontext-dev img2img pipeline."
)
if args.fp4:
if args.version == "flux.1-dev" or args.version == "flux.1-schnell":
raise ValueError("--fp4 is currently not supported for Flux.1-dev and Flux.1-schnell img2img pipelines.")
if not args.download_onnx_models:
raise ValueError("--download-onnx-models is required when using --fp4.")
kwargs_run_demo = {
"prompt": prompt,
"prompt2": prompt2,
"height": args.height,
"width": args.width,
"batch_count": args.batch_count,
"num_warmup_runs": args.num_warmup_runs,
"use_cuda_graph": args.use_cuda_graph,
"image_strength": args.image_strength,
}
# Add the appropriate image parameter based on pipeline type
if not args.version == "flux.1-kontext-dev":
kwargs_run_demo["input_image"] = args.input_image
kwargs_run_demo["control_image"] = args.control_image
return kwargs_run_demo
if __name__ == "__main__":
print("[I] Initializing Flux img2img demo using TensorRT")
args = parse_args()
_, kwargs_load_engine, _ = dd_argparse.process_pipeline_args(args)
kwargs_run_demo = process_demo_args(args)
# Initialize demo
pipeline_type = pipeline_module.PIPELINE_TYPE.IMG2IMG
if args.version == "flux.1-kontext-dev":
demo = pipeline_module.FluxKontextPipeline.FromArgs(args, pipeline_type=pipeline_type)
else:
demo = pipeline_module.FluxPipeline.FromArgs(args, pipeline_type=pipeline_type)
# Load TensorRT engines and pytorch modules
demo.load_engines(
framework_model_dir=args.framework_model_dir,
**kwargs_load_engine,
)
if args.onnx_export_only:
print("[I] ONNX export completed. Exiting...")
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.device_memory_sizes = demo.get_device_memory_sizes()
else:
_, shared_device_memory = cudart.cudaMalloc(demo.calculate_max_device_memory())
demo.activate_engines(shared_device_memory)
demo.load_resources(args.height, args.width, args.batch_size, args.seed)
# Run inference
images = demo.run(**kwargs_run_demo)
demo.teardown()
# save images
demo.save_images(kwargs_run_demo["prompt"], images, check_integrity=(args.version == "flux.1-kontext-dev"))
demo/Diffusion/demo_img2vid.py
+134
View File
@@ -0,0 +1,134 @@
#
# 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.
#
# Configure dependencies before any external imports
from demo_diffusion import deps
deps.configure("sd")
import argparse
from PIL import 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 = 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)")
parser.add_argument('--width', type=int, default=1024, help="Width of image to generate (must be multiple of 8)")
parser.add_argument('--min-guidance-scale', type=float, default=1.0, help="The minimum guidance scale. Used for the classifier free guidance with first frame")
parser.add_argument('--max-guidance-scale', type=float, default=3.0, help="The maximum guidance scale. Used for the classifier free guidance with last frame")
parser.add_argument('--denoising-steps', type=int, default=25, help="Number of denoising steps")
parser.add_argument('--num-warmup-runs', type=int, default=1, help="Number of warmup runs before benchmarking performance")
return parser.parse_args()
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 = 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:
raise ValueError(f"Input image(s) must be of type `PIL.Image.Image` or `str` (URL) but is {type(args.input_image)}")
if args.height % 8 != 0 or args.width % 8 != 0:
raise ValueError(f"Image height and width have to be divisible by 8 but are: {args.image_height} and {args.width}.")
# TODO enable BS>1
max_batch_size = 1
args.build_static_batch = True
if args.batch_size > max_batch_size:
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("Dynamic shapes not supported. Do not specify `--build-dynamic-shape`")
if args.fp8:
import torch
device_info = torch.cuda.get_device_properties(0)
version = device_info.major * 10 + device_info.minor
if version < 90: # FP8 is only supppoted on Hopper.
raise ValueError(f"Cannot apply FP8 quantization for GPU with compute capability {version / 10.0}. FP8 is only supppoted on Hopper.")
args.optimization_level = 4
print(f"[I] The default optimization level has been set to {args.optimization_level} for FP8.")
if args.quantization_level == 0.0 and args.fp8:
args.quantization_level = 3.0
print("[I] The default quantization level has been set to 3.0 for FP8.")
kwargs_init_pipeline = {
'version': args.version,
'max_batch_size': max_batch_size,
'denoising_steps': args.denoising_steps,
'scheduler': args.scheduler,
'min_guidance_scale': args.min_guidance_scale,
'max_guidance_scale': args.max_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,
'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,
'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,
'fp8': args.fp8,
'quantization_level': args.quantization_level,
}
args_run_demo = (input_image, 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
if __name__ == "__main__":
print("[I] Initializing StableDiffusion img2vid demo using TensorRT")
args = parseArgs()
kwargs_init_pipeline, kwargs_load_engine, args_run_demo = process_pipeline_args(args)
# Initialize demo
demo = pipeline_module.StableVideoDiffusionPipeline(
pipeline_type=pipeline_module.PIPELINE_TYPE.IMG2VID, **kwargs_init_pipeline
)
demo.loadEngines(
args.engine_dir,
args.framework_model_dir,
args.onnx_dir,
**kwargs_load_engine)
demo.loadResources(args.height, args.width, args.batch_size, args.seed)
# Run inference
demo.run(*args_run_demo)
demo.teardown()
demo/Diffusion/demo_stable_cascade.py
+165
View File
@@ -0,0 +1,165 @@
#
# 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.
#
# Configure dependencies before any external imports
from demo_diffusion import deps
deps.configure("sd")
import argparse
import os
import torch
from cuda.bindings import runtime as cudart
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 = 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)")
parser.add_argument('--lite', action='store_true', help="Use the Lite Version of the Stage B and Stage C models")
parser.add_argument('--prior-guidance-scale', type=float, default=4.0, help="Value of classifier-free guidance scale for the prior")
parser.add_argument('--decoder-guidance-scale', type=float, default=0.0, help="Value of classifier-free guidance scale for the decoder")
parser.add_argument('--prior-denoising-steps', type=int, default=20, help="Number of denoising steps for the prior")
parser.add_argument('--decoder-denoising-steps', type=int, default=10, help="Number of denoising steps for the decoder")
return parser.parse_args()
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 = 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 = pipeline_module.StableCascadePipeline(
pipeline_type=pipeline_module.PIPELINE_TYPE.CASCADE_DECODER,
denoising_steps=decoder_denoising_steps,
guidance_scale=decoder_guidance_scale,
lite=lite,
**kwargs,
)
def loadEngines(self, framework_model_dir, onnx_dir, engine_dir, **kwargs):
prior_suffix = "prior_lite" if self.prior.lite else "prior"
decoder_suffix = "decoder_lite" if self.decoder.lite else "decoder"
self.prior.loadEngines(
os.path.join(engine_dir, prior_suffix),
framework_model_dir,
os.path.join(onnx_dir, prior_suffix),
**kwargs)
self.decoder.loadEngines(
os.path.join(engine_dir, decoder_suffix),
framework_model_dir,
os.path.join(onnx_dir, decoder_suffix),
**kwargs)
def activateEngines(self, shared_device_memory=None):
self.prior.activateEngines(shared_device_memory)
self.decoder.activateEngines(shared_device_memory)
def loadResources(self, image_height, image_width, batch_size, seed):
self.prior.loadResources(image_height, image_width, batch_size, seed)
# Use a different seed for decoder
self.decoder.loadResources(image_height, image_width, batch_size, ((seed+1) if seed is not None else None))
def get_max_device_memory(self):
max_device_memory = self.prior.calculateMaxDeviceMemory()
max_device_memory = max(max_device_memory, self.decoder.calculateMaxDeviceMemory())
return max_device_memory
def run(self, prompt, negative_prompt, height, width, batch_size, batch_count, num_warmup_runs, use_cuda_graph):
# Process prompt
if not isinstance(prompt, list):
raise ValueError(f"`prompt` must be of type `str` list, but is {type(prompt)}")
prompt = prompt * batch_size
if not isinstance(negative_prompt, list):
raise ValueError(f"`--negative-prompt` must be of type `str` list, but is {type(negative_prompt)}")
if len(negative_prompt) == 1:
negative_prompt = negative_prompt * batch_size
num_warmup_runs = max(1, num_warmup_runs) if use_cuda_graph else num_warmup_runs
if num_warmup_runs > 0:
print("[I] Warming up ..")
for _ in range(num_warmup_runs):
latents, _ = self.prior.infer(prompt, negative_prompt, height, width, warmup=True)
latents = latents.to(torch.float16) if self.decoder.fp16 else latents
images, _ = self.decoder.infer(prompt, negative_prompt, height, width, image_embeddings=latents, warmup=True)
for _ in range(batch_count):
print("[I] Running Stable Cascade pipeline")
if self.nvtx_profile:
cudart.cudaProfilerStart()
latents, time_prior = self.prior.infer(prompt, negative_prompt, height, width, warmup=False)
latents = latents.to(torch.float16) if self.decoder.fp16 else latents
images, time_decoder = self.decoder.infer(prompt, negative_prompt, height, width, image_embeddings=latents, warmup=False)
if self.nvtx_profile:
cudart.cudaProfilerStop()
print('|-----------------|--------------|')
print('| {:^15} | {:>9.2f} ms |'.format('e2e', time_prior + time_decoder))
print('|-----------------|--------------|')
def teardown(self):
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 = dd_argparse.process_pipeline_args(args)
# Initialize demo
_ = kwargs_init_pipeline.pop('guidance_scale')
_ = kwargs_init_pipeline.pop('denoising_steps')
demo = StableCascadeDemoPipeline(
args.prior_denoising_steps,
args.decoder_denoising_steps,
args.prior_guidance_scale,
args.decoder_guidance_scale,
args.lite,
**kwargs_init_pipeline
)
# Load TensorRT engines and pytorch modules
demo.loadEngines(
args.framework_model_dir,
args.onnx_dir,
args.engine_dir,
**kwargs_load_engine,
)
# Load resources
_, shared_device_memory = cudart.cudaMalloc(demo.get_max_device_memory())
demo.activateEngines(shared_device_memory)
demo.loadResources(args.height, args.width, args.batch_size, args.seed)
# Run inference
demo.run(*args_run_demo)
demo.teardown()
demo/Diffusion/demo_txt2image_cosmos.py
+151
View File
@@ -0,0 +1,151 @@
#
# 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.
#
# Configure dependencies before any external imports
from demo_diffusion import deps
deps.configure("cosmos")
import argparse
from cuda.bindings import runtime as cudart
from demo_diffusion import dd_argparse
from demo_diffusion import pipeline as pipeline_module
def parse_args():
parser = argparse.ArgumentParser(description="Options for Cosmos text2image Demo", conflict_handler="resolve")
parser = dd_argparse.add_arguments(parser)
parser.add_argument(
"--version",
type=str,
default="cosmos-predict2-2b-text2image",
choices=("cosmos-predict2-2b-text2image", "cosmos-predict2-14b-text2image"),
help="Version of Cosmos",
)
parser.add_argument(
"--height",
type=int,
default=768,
help="Height of image to generate (must be multiple of 8)",
)
parser.add_argument(
"--width",
type=int,
default=1360,
help="Width of image to generate (must be multiple of 8)",
)
parser.add_argument("--denoising-steps", type=int, default=35, help="Number of denoising steps")
parser.add_argument(
"--guidance-scale",
type=float,
default=7.0,
help="Value of classifier-free guidance scale (must be greater than 1)",
)
parser.add_argument(
"--num-images-per-prompt", type=int, default=1, help="The number of images to generate per prompt."
)
parser.add_argument(
"--max_sequence_length",
type=int,
default=512,
help="Maximum sequence length to use with the prompt.",
)
parser.add_argument(
"--t5-ws-percentage",
type=int,
default=None,
help="Set runtime weight streaming budget as the percentage of the size of streamable weights for the T5 model. This argument only takes effect when --ws is set. 0 streams the most weights and 100 or None streams no weights. ",
)
parser.add_argument(
"--transformer-ws-percentage",
type=int,
default=None,
help="Set runtime weight streaming budget as the percentage of the size of streamable weights for the transformer model. This argument only takes effect when --ws is set. 0 streams the most weights and 100 or None streams no weights.",
)
parser.add_argument(
"--bf16",
action="store_true",
default=True,
help="Use bfloat16 precision by default.",
)
return parser.parse_args()
def process_demo_args(args):
batch_size = args.batch_size
prompt = args.prompt
negative_prompt = args.negative_prompt
# Process input args
if not isinstance(prompt, list):
raise ValueError(f"`prompt` must be of type `str` list, but is {type(prompt)}")
prompt = prompt * batch_size
if not isinstance(negative_prompt, list):
raise ValueError(f"`negative_prompt` must be of type `str` list, but is {type(negative_prompt)}")
negative_prompt = negative_prompt * batch_size
kwargs_run_demo = {
"prompt": prompt,
"negative_prompt": negative_prompt,
"height": args.height,
"width": args.width,
"batch_count": args.batch_count,
"num_warmup_runs": args.num_warmup_runs,
"use_cuda_graph": args.use_cuda_graph,
"num_images_per_prompt": args.num_images_per_prompt,
}
return kwargs_run_demo
if __name__ == "__main__":
print("[I] Initializing Cosmos text2image demo using TensorRT")
args = parse_args()
_, kwargs_load_engine, _ = dd_argparse.process_pipeline_args(args)
kwargs_run_demo = process_demo_args(args)
# Initialize demo
demo = pipeline_module.CosmosPipeline.FromArgs(args, pipeline_type=pipeline_module.PIPELINE_TYPE.TXT2IMG)
# Load TensorRT engines and pytorch modules
demo.load_engines(
framework_model_dir=args.framework_model_dir,
**kwargs_load_engine,
)
if args.onnx_export_only:
print("[I] ONNX export completed. Exiting...")
demo.teardown()
exit(0)
# In low-vram mode we allocate the required device memory individually before each model is run.
if demo.low_vram:
demo.device_memory_sizes = demo.get_device_memory_sizes()
else:
_, shared_device_memory = cudart.cudaMalloc(demo.calculate_max_device_memory())
demo.activate_engines(shared_device_memory)
demo.load_resources(args.height, args.width, args.batch_size, args.seed)
# Run inference
images = demo.run(**kwargs_run_demo)
demo.teardown()
# save images
demo.save_images(kwargs_run_demo["prompt"], images, check_integrity=True)
demo/Diffusion/demo_txt2img.py
+61
View File
@@ -0,0 +1,61 @@
#
# 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.
#
# Configure dependencies before any external imports
from demo_diffusion import deps
deps.configure("sd")
import argparse
from cuda.bindings import runtime as cudart
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 = 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 = dd_argparse.process_pipeline_args(args)
# Initialize demo
demo = pipeline_module.StableDiffusionPipeline(
pipeline_type=pipeline_module.PIPELINE_TYPE.TXT2IMG, **kwargs_init_pipeline
)
# Load TensorRT engines and pytorch modules
demo.loadEngines(
args.engine_dir,
args.framework_model_dir,
args.onnx_dir,
**kwargs_load_engine)
# Load resources
_, shared_device_memory = cudart.cudaMalloc(demo.calculateMaxDeviceMemory())
demo.activateEngines(shared_device_memory)
demo.loadResources(args.height, args.width, args.batch_size, args.seed)
# Run inference
demo.run(*args_run_demo)
demo.teardown()
demo/Diffusion/demo_txt2img_flux.py
+171
View File
@@ -0,0 +1,171 @@
#
# 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.
#
# Configure dependencies before any external imports
from demo_diffusion import deps
deps.configure("flux")
import argparse
from cuda.bindings import runtime as 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 = dd_argparse.add_arguments(parser)
parser.add_argument(
"--version",
type=str,
default="flux.1-dev",
choices=("flux.1-dev", "flux.1-schnell"),
help="Version of Flux",
)
parser.add_argument(
"--prompt2",
default=None,
nargs="*",
help="Text prompt(s) to be sent to the T5 tokenizer and text encoder. If not defined, prompt will be used instead",
)
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)",
)
parser.add_argument(
"--denoising-steps", type=int, default=50, help="Number of denoising steps"
)
parser.add_argument(
"--guidance-scale",
type=float,
default=3.5,
help="Value of classifier-free guidance scale (must be greater than 1)",
)
parser.add_argument(
"--max_sequence_length",
type=int,
help="Maximum sequence length to use with the prompt. Can be up to 512 for the dev and 256 for the schnell variant.",
)
parser.add_argument(
"--t5-ws-percentage",
type=int,
default=None,
help="Set runtime weight streaming budget as the percentage of the size of streamable weights for the T5 model. This argument only takes effect when --ws is set. 0 streams the most weights and 100 or None streams no weights. ",
)
parser.add_argument(
"--transformer-ws-percentage",
type=int,
default=None,
help="Set runtime weight streaming budget as the percentage of the size of streamable weights for the transformer model. This argument only takes effect when --ws is set. 0 streams the most weights and 100 or None streams no weights."
)
return parser.parse_args()
def process_demo_args(args):
batch_size = args.batch_size
prompt = args.prompt
# If prompt2 is not defined, use prompt instead
prompt2 = args.prompt2 or prompt
# Process input args
if not isinstance(prompt, list):
raise ValueError(f"`prompt` must be of type `str` list, but is {type(prompt)}")
prompt = prompt * batch_size
if not isinstance(prompt2, list):
raise ValueError(
f"`prompt2` must be of type `str` list, but is {type(prompt2)}"
)
if len(prompt2) == 1:
prompt2 = prompt2 * batch_size
max_seq_supported_by_model = {
"flux.1-schnell": 256,
"flux.1-dev": 512,
}[args.version]
if args.max_sequence_length is not None:
if args.max_sequence_length > max_seq_supported_by_model:
raise ValueError(
f"For {args.version}, `max_sequence_length` cannot be greater than {max_seq_supported_by_model} but is {args.max_sequence_length}"
)
else:
args.max_sequence_length = max_seq_supported_by_model
kwargs_run_demo = {
"prompt": prompt,
"prompt2": prompt2,
"height": args.height,
"width": args.width,
"batch_count": args.batch_count,
"num_warmup_runs": args.num_warmup_runs,
"use_cuda_graph": args.use_cuda_graph,
}
return kwargs_run_demo
if __name__ == "__main__":
print("[I] Initializing Flux txt2img demo using TensorRT")
args = parse_args()
_, kwargs_load_engine, _ = dd_argparse.process_pipeline_args(args)
kwargs_run_demo = process_demo_args(args)
# Initialize demo
pipeline_type = pipeline_module.PIPELINE_TYPE.TXT2IMG
demo = pipeline_module.FluxPipeline.FromArgs(args, pipeline_type=pipeline_type)
# Load TensorRT engines and pytorch modules
demo.load_engines(
framework_model_dir=args.framework_model_dir,
**kwargs_load_engine,
)
if args.onnx_export_only:
print("[I] ONNX export completed. Exiting...")
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.device_memory_sizes = demo.get_device_memory_sizes()
else:
_, shared_device_memory = cudart.cudaMalloc(demo.calculate_max_device_memory())
demo.activate_engines(shared_device_memory)
demo.load_resources(args.height, args.width, args.batch_size, args.seed)
# Run inference
images = demo.run(**kwargs_run_demo)
demo.teardown()
# save images
demo.save_images(kwargs_run_demo["prompt"], images)
demo/Diffusion/demo_txt2img_sd3.py
+130
View File
@@ -0,0 +1,130 @@
#
# 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.
#
# Configure dependencies before any external imports
from demo_diffusion import deps
deps.configure("sd")
import argparse
from cuda.bindings import runtime as 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 = 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)")
parser.add_argument('--shift', type=int, default=1.0, help="Shift parameter for SD3")
parser.add_argument('--cfg-scale', type=int, default=5, help="CFG Scale for SD3")
parser.add_argument('--denoising-steps', type=int, default=50, help="Number of denoising steps")
parser.add_argument('--denoising-percentage', type=float, default=0.6, help="Percentage of denoising steps to run. This parameter is only used if input-image is provided")
parser.add_argument('--input-image', type=str, default="", help="Path to the input image")
return parser.parse_args()
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`"
)
input_image = None
if args.input_image:
input_image = Image.open(args.input_image)
image_width, image_height = input_image.size
if image_height != args.height or image_width != args.width:
print(f"[I] Resizing input_image to {args.height}x{args.width}")
input_image = input_image.resize((args.width, args.height), Image.LANCZOS)
image_height, image_width = args.height, args.width
input_image = preprocess_image_sd3(input_image)
kwargs_init_pipeline = {
'version': args.version,
'max_batch_size': max_batch_size,
'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,
'torch_inference': args.torch_inference,
'shift': args.shift,
'cfg_scale': args.cfg_scale,
'denoising_steps': args.denoising_steps,
'denoising_percentage': args.denoising_percentage,
'input_image': input_image
}
kwargs_load_engine = {
'onnx_opset': args.onnx_opset,
'opt_batch_size': args.batch_size,
'opt_image_height': args.height,
'opt_image_width': args.width,
'static_batch': args.build_static_batch,
'static_shape': not args.build_dynamic_shape,
'enable_all_tactics': args.build_all_tactics,
'timing_cache': args.timing_cache,
}
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
if __name__ == "__main__":
print("[I] Initializing Stable Diffusion 3 demo using TensorRT")
args = parseArgs()
kwargs_init_pipeline, kwargs_load_engine, args_run_demo = process_pipeline_args(args)
# Initialize demo
demo = pipeline_module.StableDiffusion3Pipeline(
pipeline_type=pipeline_module.PIPELINE_TYPE.TXT2IMG, **kwargs_init_pipeline
)
# Load TensorRT engines and pytorch modules
demo.loadEngines(
args.engine_dir,
args.framework_model_dir,
args.onnx_dir,
**kwargs_load_engine)
# Load resources
_, shared_device_memory = cudart.cudaMalloc(demo.calculateMaxDeviceMemory())
demo.activateEngines(shared_device_memory)
demo.loadResources(args.height, args.width, args.batch_size, args.seed)
# Run inference
demo.run(*args_run_demo)
demo.teardown()
demo/Diffusion/demo_txt2img_sd35.py
+135
View File
@@ -0,0 +1,135 @@
#
# 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.
#
# Configure dependencies before any external imports
from demo_diffusion import deps
deps.configure("sd")
import argparse
from cuda.bindings import runtime as cudart
from demo_diffusion import dd_argparse
from demo_diffusion import pipeline as pipeline_module
def parseArgs():
# Stable Diffusion 3.5 configuration
parser = argparse.ArgumentParser(
description="Options for Stable Diffusion 3.5 Txt2Img Demo", conflict_handler="resolve"
)
parser = dd_argparse.add_arguments(parser)
parser.add_argument(
"--version",
type=str,
default="3.5-medium",
choices={"3.5-medium", "3.5-large"},
help="Version of Stable Diffusion 3.5",
)
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)")
parser.add_argument(
"--guidance-scale",
type=float,
default=7.0,
help="Value of classifier-free guidance scale (must be greater than 1)",
)
parser.add_argument(
"--max-sequence-length",
type=int,
default=256,
help="Maximum sequence length to use with the prompt.",
)
parser.add_argument("--denoising-steps", type=int, default=50, help="Number of denoising steps")
return parser.parse_args()
def process_demo_args(args):
batch_size = args.batch_size
prompt = args.prompt
negative_prompt = args.negative_prompt
# Process prompt
if not isinstance(prompt, list):
raise ValueError(f"`prompt` must be of type `str` list, but is {type(prompt)}")
prompt = prompt * batch_size
if not isinstance(negative_prompt, list):
raise ValueError(f"`--negative-prompt` must be of type `str` list, but is {type(negative_prompt)}")
if len(negative_prompt) == 1:
negative_prompt = negative_prompt * batch_size
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`"
)
kwargs_run_demo = {
"prompt": prompt,
"negative_prompt": negative_prompt,
"height": args.height,
"width": args.width,
"batch_count": args.batch_count,
"num_warmup_runs": args.num_warmup_runs,
"use_cuda_graph": args.use_cuda_graph,
}
return kwargs_run_demo
if __name__ == "__main__":
print("[I] Initializing Stable Diffusion 3.5 demo using TensorRT")
args = parseArgs()
_, kwargs_load_engine, _ = dd_argparse.process_pipeline_args(args)
kwargs_run_demo = process_demo_args(args)
# Initialize demo
demo = pipeline_module.StableDiffusion35Pipeline.FromArgs(args, pipeline_type=pipeline_module.PIPELINE_TYPE.TXT2IMG)
# Load TensorRT engines and pytorch modules
demo.load_engines(
framework_model_dir=args.framework_model_dir,
**kwargs_load_engine,
)
if demo.low_vram:
demo.device_memory_sizes = demo.get_device_memory_sizes()
else:
_, shared_device_memory = cudart.cudaMalloc(demo.calculate_max_device_memory())
demo.activate_engines(shared_device_memory)
# Load resources
demo.load_resources(
image_height=args.height,
image_width=args.width,
batch_size=args.batch_size,
seed=args.seed,
)
# Run inference
demo.run(**kwargs_run_demo)
demo.teardown()
demo/Diffusion/demo_txt2img_xl.py
+160
View File
@@ -0,0 +1,160 @@
#
# 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.
#
# Configure dependencies before any external imports
from demo_diffusion import deps
deps.configure("sd")
import argparse
from cuda.bindings import runtime as cudart
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 = 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)")
parser.add_argument('--num-warmup-runs', type=int, default=1, help="Number of warmup runs before benchmarking performance")
parser.add_argument('--guidance-scale', type=float, default=5.0, help="Value of classifier-free guidance scale (must be greater than 1)")
parser.add_argument('--enable-refiner', action='store_true', help="Enable SDXL-Refiner model")
parser.add_argument('--image-strength', type=float, default=0.3, help="Strength of transformation applied to input_image (must be between 0 and 1)")
parser.add_argument('--onnx-refiner-dir', default='onnx_xl_refiner', help="Directory for SDXL-Refiner ONNX models")
parser.add_argument('--engine-refiner-dir', default='engine_xl_refiner', help="Directory for SDXL-Refiner TensorRT engines")
return parser.parse_args()
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 = pipeline_module.StableDiffusionPipeline(
pipeline_type=pipeline_module.PIPELINE_TYPE.XL_BASE,
vae_scaling_factor=vae_scaling_factor,
return_latents=self.enable_refiner,
**kwargs,
)
if self.enable_refiner:
self.refiner = pipeline_module.StableDiffusionPipeline(
pipeline_type=pipeline_module.PIPELINE_TYPE.XL_REFINER,
vae_scaling_factor=vae_scaling_factor,
return_latents=False,
**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)
if self.enable_refiner:
self.refiner.loadEngines(engine_refiner_dir, framework_model_dir, onnx_refiner_dir, **kwargs)
def activateEngines(self, shared_device_memory=None):
self.base.activateEngines(shared_device_memory)
if self.enable_refiner:
self.refiner.activateEngines(shared_device_memory)
def loadResources(self, image_height, image_width, batch_size, seed):
self.base.loadResources(image_height, image_width, batch_size, seed)
if self.enable_refiner:
# Use a different seed for refiner - we arbitrarily use base seed+1, if specified.
self.refiner.loadResources(image_height, image_width, batch_size, ((seed+1) if seed is not None else None))
def get_max_device_memory(self):
max_device_memory = self.base.calculateMaxDeviceMemory()
if self.enable_refiner:
max_device_memory = max(max_device_memory, self.refiner.calculateMaxDeviceMemory())
return max_device_memory
def run(self, prompt, negative_prompt, height, width, batch_size, batch_count, num_warmup_runs, use_cuda_graph, **kwargs_infer_refiner):
# Process prompt
if not isinstance(prompt, list):
raise ValueError(f"`prompt` must be of type `str` list, but is {type(prompt)}")
prompt = prompt * batch_size
if not isinstance(negative_prompt, list):
raise ValueError(f"`--negative-prompt` must be of type `str` list, but is {type(negative_prompt)}")
if len(negative_prompt) == 1:
negative_prompt = negative_prompt * batch_size
num_warmup_runs = max(1, num_warmup_runs) if use_cuda_graph else num_warmup_runs
if num_warmup_runs > 0:
print("[I] Warming up ..")
for _ in range(num_warmup_runs):
images, _ = self.base.infer(prompt, negative_prompt, height, width, warmup=True)
if args.enable_refiner:
images, _ = self.refiner.infer(prompt, negative_prompt, height, width, input_image=images, warmup=True, **kwargs_infer_refiner)
ret = []
for _ in range(batch_count):
print("[I] Running StableDiffusionXL pipeline")
if self.nvtx_profile:
cudart.cudaProfilerStart()
latents, time_base = self.base.infer(prompt, negative_prompt, height, width, warmup=False)
if self.enable_refiner:
images, time_refiner = self.refiner.infer(prompt, negative_prompt, height, width, input_image=latents, warmup=False, **kwargs_infer_refiner)
ret.append(images)
else:
ret.append(latents)
if self.nvtx_profile:
cudart.cudaProfilerStop()
if self.enable_refiner:
print('|-----------------|--------------|')
print('| {:^15} | {:>9.2f} ms |'.format('e2e', time_base + time_refiner))
print('|-----------------|--------------|')
return ret
def teardown(self):
self.base.teardown()
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 = dd_argparse.process_pipeline_args(args)
# Initialize demo
demo = StableDiffusionXLPipeline(vae_scaling_factor=0.13025, enable_refiner=args.enable_refiner, **kwargs_init_pipeline)
# Load TensorRT engines and pytorch modules
kwargs_load_refiner = {'onnx_refiner_dir': args.onnx_refiner_dir, 'engine_refiner_dir': args.engine_refiner_dir} if args.enable_refiner else {}
demo.loadEngines(
args.framework_model_dir,
args.onnx_dir,
args.engine_dir,
**kwargs_load_refiner,
**kwargs_load_engine)
# Load resources
_, shared_device_memory = cudart.cudaMalloc(demo.get_max_device_memory())
demo.activateEngines(shared_device_memory)
demo.loadResources(args.height, args.width, args.batch_size, args.seed)
# Run inference
kwargs_infer_refiner = {'image_strength': args.image_strength} if args.enable_refiner else {}
demo.run(*args_run_demo, **kwargs_infer_refiner)
demo.teardown()
demo/Diffusion/demo_txt2vid_wan.py
+128
View File
@@ -0,0 +1,128 @@
#
# 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.
#
# Configure dependencies before any external imports
from demo_diffusion import deps
deps.configure("cosmos")
import argparse
from cuda.bindings import runtime as cudart
from demo_diffusion import dd_argparse
from demo_diffusion import pipeline as pipeline_module
def parseArgs():
parser = argparse.ArgumentParser(description="Options for Wan 2.2 Txt2Vid Demo", conflict_handler='resolve')
parser = dd_argparse.add_arguments(parser)
parser.add_argument('--version', type=str, default="wan2.2-t2v-a14b", help="Version of Wan")
parser.add_argument('--guidance-scale', type=float, default=4.0, help="Guidance scale for high-noise stage (Wan default: 4.0)")
parser.add_argument('--guidance-scale-2', type=float, default=3.0, help="Guidance scale for low-noise stage (Wan default: 3.0)")
parser.add_argument('--boundary-ratio', type=float, default=0.875, help="Boundary ratio for two-stage denoising (default: 0.875)")
parser.add_argument('--denoising-steps', type=int, default=40, help="Number of denoising steps (Wan default: 40)")
parser.add_argument('--num-warmup-runs', type=int, default=1, help="Number of warmup runs before benchmarking")
parser.add_argument(
'--negative-prompt',
nargs='*',
default= (
"vivid colors, overexposed, static, blurry details, subtitles, style, "
"work of art, painting, picture, still, overall grayish, worst quality, "
"low quality, JPEG artifacts, ugly, deformed, extra fingers, poorly drawn hands, "
"poorly drawn face, deformed, disfigured, deformed limbs, fused fingers, "
"static image, cluttered background, three legs, many people in the background, "
"walking backwards"
),
help="Negative prompt (Wan team default, English translation)"
)
parser.add_argument(
"--onnx-opset",
type=int,
default=23,
choices=range(7, 24),
help="Select ONNX opset version to target for exported models",
)
return parser.parse_args()
def process_demo_args(args):
args.height = 720
args.width = 1280
args.num_frames = 81
args.max_sequence_length = 512
# require static batch = 1
if args.batch_size > 1:
raise ValueError(f"Batch size {args.batch_size} is larger than allowed (max=1 for Wan).")
args.batch_size = 1
args.build_static_batch = True
args.build_dynamic_shape = False
print(f"[I] Building Wan 2.2 T2V with fixed resolution: {args.height}×{args.width}, {args.num_frames} frames")
negative_prompt = args.negative_prompt
if isinstance(negative_prompt, list):
negative_prompt = ' '.join(negative_prompt) if negative_prompt else ""
kwargs_run_demo = {
'prompt': args.prompt,
'height': args.height,
'width': args.width,
'num_frames': args.num_frames,
'batch_size': args.batch_size,
'batch_count': args.batch_count,
'num_warmup_runs': args.num_warmup_runs,
'use_cuda_graph': args.use_cuda_graph,
'negative_prompt': negative_prompt,
'num_inference_steps': args.denoising_steps,
}
return kwargs_run_demo
if __name__ == "__main__":
print("[I] Initializing Wan 2.2 txt2vid demo using TensorRT")
args = parseArgs()
kwargs_run_demo = process_demo_args(args)
_, kwargs_load_engine, _ = dd_argparse.process_pipeline_args(args)
# Initialize demo
demo = pipeline_module.WanPipeline.FromArgs(args, pipeline_type=pipeline_module.PIPELINE_TYPE.TXT2VID)
# Load TensorRT engines and pytorch modules
demo.load_engines(
framework_model_dir=args.framework_model_dir,
**kwargs_load_engine,
)
# In low-vram mode we allocate the required device memory individually before each model is run
if args.low_vram:
demo.device_memory_sizes = demo.get_device_memory_sizes()
else:
_, shared_device_memory = cudart.cudaMalloc(demo.calculate_max_device_memory())
demo.activate_engines(shared_device_memory)
# Load resources
demo.load_resources(args.height, args.width, args.batch_size, args.seed)
# Run inference
demo.run(**kwargs_run_demo)
demo.teardown()
demo/Diffusion/demo_vid2world_cosmos.py
+174
View File
@@ -0,0 +1,174 @@
#
# 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.
#
# Configure dependencies before any external imports
from demo_diffusion import deps
deps.configure("cosmos")
import argparse
from cuda.bindings import runtime as cudart
from diffusers.utils import load_image, load_video
from demo_diffusion import dd_argparse
from demo_diffusion import pipeline as pipeline_module
def parse_args():
parser = argparse.ArgumentParser(description="Options for Cosmos video2world Demo", conflict_handler="resolve")
parser = dd_argparse.add_arguments(parser)
parser.add_argument(
"--version",
type=str,
default="cosmos-predict2-2b-video2world",
choices=("cosmos-predict2-2b-video2world", "cosmos-predict2-14b-video2world"),
help="Version of Cosmos",
)
parser.add_argument('--input-image', type=str, default=None, help="Path to the input image")
parser.add_argument('--input-video', type=str, default=None, help="Path to the input video")
parser.add_argument(
"--height",
type=int,
default=704,
help="Height of image to generate (must be multiple of 8)",
)
parser.add_argument(
"--width",
type=int,
default=1280,
help="Width of image to generate (must be multiple of 8)",
)
parser.add_argument("--denoising-steps", type=int, default=35, help="Number of denoising steps")
parser.add_argument(
"--guidance-scale",
type=float,
default=7.0,
help="Value of classifier-free guidance scale (must be greater than 1)",
)
parser.add_argument("--num-frames", type=int, default=93, help="The number of frames in the generated video.")
parser.add_argument("--fps", type=int, default=16, help="The frames per second of the generated video.")
parser.add_argument("--num-videos-per-prompt", type=int, default=1, help="The number of videos to generate per prompt.")
parser.add_argument(
"--max_sequence_length",
type=int,
default=512,
help="Maximum sequence length to use with the prompt.",
)
parser.add_argument(
"--t5-ws-percentage",
type=int,
default=None,
help="Set runtime weight streaming budget as the percentage of the size of streamable weights for the T5 model. This argument only takes effect when --ws is set. 0 streams the most weights and 100 or None streams no weights. ",
)
parser.add_argument(
"--transformer-ws-percentage",
type=int,
default=None,
help="Set runtime weight streaming budget as the percentage of the size of streamable weights for the transformer model. This argument only takes effect when --ws is set. 0 streams the most weights and 100 or None streams no weights.",
)
parser.add_argument(
"--bf16",
action="store_true",
default=True,
help="Use bfloat16 precision by default.",
)
return parser.parse_args()
def process_demo_args(args):
batch_size = args.batch_size
prompt = args.prompt
negative_prompt = args.negative_prompt
# Process input args
if not isinstance(prompt, list):
raise ValueError(f"`prompt` must be of type `str` list, but is {type(prompt)}")
prompt = prompt * batch_size
if not isinstance(negative_prompt, list):
raise ValueError(f"`negative_prompt` must be of type `str` list, but is {type(negative_prompt)}")
negative_prompt = negative_prompt * batch_size
# process input image and input video
if args.input_image and args.input_video:
raise ValueError("Only one of --input-image or --input-video can be provided")
if args.input_image:
args.input_image = load_image(args.input_image)
elif args.input_video:
args.input_video = load_video(args.input_video)
else:
# load default image
args.input_image = load_image("https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/yellow-scrubber.png")
kwargs_run_demo = {
"prompt": prompt,
"negative_prompt": negative_prompt,
"height": args.height,
"width": args.width,
"batch_count": args.batch_count,
"num_warmup_runs": args.num_warmup_runs,
"use_cuda_graph": args.use_cuda_graph,
"num_frames": args.num_frames,
"fps": args.fps,
"input_image": args.input_image,
"input_video": args.input_video,
"num_videos_per_prompt": args.num_videos_per_prompt,
}
return kwargs_run_demo
if __name__ == "__main__":
print("[I] Initializing Cosmos video2world demo using TensorRT")
args = parse_args()
# Enforce torch-inference is enabled
if not args.torch_inference:
print("[W] The video2world demo only supports the PyTorch backend. Enabling torch-inference with 'eager' mode.")
args.torch_inference = "eager"
_, kwargs_load_engine, _ = dd_argparse.process_pipeline_args(args)
kwargs_run_demo = process_demo_args(args)
# Initialize demo
demo = pipeline_module.CosmosPipeline.FromArgs(args, pipeline_type=pipeline_module.PIPELINE_TYPE.VIDEO2WORLD)
# Load TensorRT engines and pytorch modules
demo.load_engines(
framework_model_dir=args.framework_model_dir,
**kwargs_load_engine,
)
if args.onnx_export_only:
print("[I] ONNX export completed. Exiting...")
demo.teardown()
exit(0)
# In low-vram mode we allocate the required device memory individually before each model is run.
if demo.low_vram:
demo.device_memory_sizes = demo.get_device_memory_sizes()
else:
_, shared_device_memory = cudart.cudaMalloc(demo.calculate_max_device_memory())
demo.activate_engines(shared_device_memory)
demo.load_resources(args.height, args.width, args.batch_size, args.seed)
# Run inference
videos = demo.run(**kwargs_run_demo)
demo.teardown()
# save video
demo.save_video(kwargs_run_demo["prompt"], videos, check_integrity=True)
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 | Additional features | Hub | Restrictions |
|------------------------|----------------------------------------------------------------------------------------------------------------------------------------------------------------------|-------------------------------------------------------------|-----------------------------------------------|--------------------------|--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------|----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------|
| Stable Diffusion | 1.4 | <ul><li>Text-to-image</li><li>Image-to-image</li></ul> | FP16, FP8, INT8 * | N/A | [CompVis/stable-diffusion-v1-4](https://huggingface.co/CompVis/stable-diffusion-v1-4) | |
| Stable Diffusion | dreamshaper-7 | <ul><li>Text-to-image</li><li>Image-to-image</li></ul> | FP16 | N/A | [Lykon/dreamshaper-7](https://huggingface.co/Lykon/dreamshaper-7) |
| 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 * | LoRA (FP16, BF16, FP8) | [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 | N/A | [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 | N/A | [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 | N/A | [stabilityai/stable-diffusion-3-medium](https://huggingface.co/stabilityai/stable-diffusion-3-medium) |
| Stable Diffusion | [3.5-medium](../README.md#generate-an-image-guided-by-a-text-prompt-using-stable-diffusion-3) | <ul><li>Text-to-image</li></ul> | FP16, BF16 | N/A | [stabilityai/stable-diffusion-3-medium](https://huggingface.co/stabilityai/stable-diffusion-3-medium) |
| Stable Diffusion | [3.5-large](../README.md#generate-an-image-guided-by-a-text-prompt-using-stable-diffusion-3) | <ul><li>Text-to-image</li></ul> | FP16, BF16, FP8 | N/A | [stabilityai/stable-diffusion-3-large](https://huggingface.co/stabilityai/stable-diffusion-3-large) |
| ControlNet | [1.4](../README.md#generate-an-image-with-controlnet-guided-by-images-and-text-prompts) | <ul><li>Image-to-image</li></ul> | FP16 | N/A | <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](../README.md#generate-an-image-with-stable-diffusion-xl-guided-by-a-single-text-prompt) | <ul><li>Image-to-image</li></ul> | FP16, FP8 | N/A | [stabilityai/controlnet-canny-sdxl-1.0](https://huggingface.co/diffusers/controlnet-canny-sdxl-1.0) |
| ControlNet | [3.5-large](../README.md#generate-an-image-with-stable-diffusion-v35-large-with-controlnet-guided-by-an-image-and-a-text-prompt) | <ul><li>Image-to-image</li></ul> | FP16, BF16, FP8 (canny and depth only) | N/A | <ul><li>[stabilityai/stable-diffusion-3.5-large-controlnet-canny](https://huggingface.co/stabilityai/stable-diffusion-3.5-large-controlnet-canny)</li><li>[stabilityai/stable-diffusion-3.5-large-controlnet-depth](https://huggingface.co/stabilityai/stable-diffusion-3.5-large-controlnet-depth)</li><li>[stabilityai/stable-diffusion-3.5-large-controlnet-blur](https://huggingface.co/stabilityai/stable-diffusion-3.5-large-controlnet-blur)</li></ul> |
| 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 | N/A | [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 | N/A | <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 * | LoRA (FP16, BF16, FP8) | [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 * | LoRA (FP16, BF16, FP8) | [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 | N/A | [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 | N/A | [black-forest-labs/FLUX.1-Depth-dev](https://huggingface.co/black-forest-labs/FLUX.1-Depth-dev) |
| Flux | [1-Kontext-Dev](../README.md#5-edit-an-image-using-flux-kontext) | <ul><li>Image-to-image</li></ul> | BF16, FP8, FP4 | N/A | [black-forest-labs/FLUX.1-Kontext-dev](https://huggingface.co/black-forest-labs/FLUX.1-Kontext-dev) |
| Cosmos | [cosmos-predict2-2b-text2image](../README.md#1-generate-an-image-from-a-text-prompt-1), cosmos-predict2-14b-text2image | <ul><li>Text-to-image</li></ul> | BF16 | N/A | <ul><li>[nvidia/Cosmos-Predict2-2B-Text2Image](https://huggingface.co/nvidia/Cosmos-Predict2-2B-Text2Image)</li><li>[nvidia/Cosmos-Predict2-14B-Text2Image](https://huggingface.co/nvidia/Cosmos-Predict2-14B-Text2Image) |
| Cosmos | [cosmos-predict2-2b-video2world](../README.md#2-generate-a-video-guided-by-an-initial-video-conditioning-and-a-text-prompt), cosmos-predict2-14b-video2world | <ul><li>Video-to-World</li></ul> | BF16 | N/A | <ul><li>[nvidia/Cosmos-Predict2-2B-Video2World](https://huggingface.co/nvidia/Cosmos-Predict2-2B-Video2World)</li><li>[nvidia/Cosmos-Predict2-14B-Video2World](https://huggingface.co/nvidia/Cosmos-Predict2-14B-Video2World) |
| Wan | [2.2-T2V-A14B](../README.md#generate-a-video-from-a-text-prompt-using-wan) | <ul><li>Text-to-video</li></ul> | BF16 | N/A | [wan-ai/wan2.2-t2v-a14b](https://huggingface.co/Wan-AI/Wan2.2-T2V-A14B-Diffusers) |
*Note: Only the text2image pipelines support FP4/FP8/INT8 quantization. The image2image pipelines don't support quantization.
demo/Diffusion/requirements.txt
Executable
+24
View File
@@ -0,0 +1,24 @@
apex==0.9.10dev
accelerate==1.2.1
colored
controlnet_aux==0.0.6
cuda-python
diffusers==0.35.2
git+https://github.com/black-forest-labs/flux.git
ftfy
matplotlib
nvtx
onnx==1.18.0
onnxscript==0.5.4
opencv-python-headless==4.8.0.74
scipy
transformers==4.52.4
--extra-index-url https://pypi.nvidia.com
nvidia-modelopt[torch,onnx]==0.31.0
onnx-graphsurgeon==0.5.2
peft==0.17.0
polygraphy==0.49.22
sentencepiece
numpy==1.26.4
imageio-ffmpeg
demo/Diffusion/setup.py
+578
View File
@@ -0,0 +1,578 @@
#!/usr/bin/env python3
#
# 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.
#
"""
Setup script for TensorRT Diffusion dependencies.
Usage:
python setup.py [GROUP] [OPTIONS]
Groups:
sd = SD 1.4, SDXL, SD3, SD3.5, SVD, Stable Cascade (Stability AI)
flux = Flux models (Black Forest Labs)
cosmos = Cosmos models (NVIDIA), Wan2.2 T2V
all = Install all groups (default)
Options:
--skip-tensorrt Skip TensorRT upgrade/installation
--force Force reinstallation even if already installed
--deps-root DIR Root directory for dependencies
-q, --quiet Suppress informational output (errors are always shown)
Examples:
python setup.py # Install all
python setup.py sd # Install SD family only
python setup.py flux # Install Flux only
python setup.py cosmos # Install Cosmos only
python setup.py --skip-tensorrt # Install all, skip TensorRT upgrade
python setup.py flux --skip-tensorrt # Install Flux only, skip TensorRT upgrade
python setup.py all --quiet # Install all, minimal output
"""
import argparse
import os
import re
import shutil
import subprocess
import sys
import tempfile
from pathlib import Path
# Global quiet flag (set by parse_args)
_quiet = False
def log(msg: str = ""):
"""Print a message unless in quiet mode."""
if not _quiet:
print(msg)
# Group descriptions
GROUP_DESCRIPTIONS = {
"sd": "SD family (SD 1.4, SDXL, SD3, SD3.5, SVD, Stable Cascade)",
"flux": "Flux family (Black Forest Labs)",
"cosmos": "Cosmos family (NVIDIA), Wan2.2 T2V",
}
# Valid dependency groups
VALID_GROUPS = set(GROUP_DESCRIPTIONS.keys())
# Default installation root
# Can be overridden by:
# 1. Environment variable: TENSORRT_DIFFUSION_DEPS_ROOT
# 2. Command-line argument: --deps-root
DEFAULT_DEPS_ROOT = os.environ.get("TENSORRT_DIFFUSION_DEPS_ROOT", "/workspace/deps")
# Marker file to indicate successful installation
# This file is created after all packages are successfully installed
# Prevents treating incomplete installations as complete
INSTALL_COMPLETE_MARKER = ".install_complete"
def _normalize_package_name(name: str) -> str:
"""Normalize a package name per PEP 503 (e.g. 'Nvidia_ModelOpt' -> 'nvidia-modelopt')."""
return re.sub(r"[-_.]+", "-", name).lower()
# Regex to extract the package name from a PEP 508 requirement string.
# Matches the leading identifier before any extras, version specifier, or URL marker.
# "nvidia-modelopt[torch,onnx]==0.40.0" -> "nvidia-modelopt"
# "flux @ git+https://..." -> "flux"
_REQ_NAME_RE = re.compile(r"^([A-Za-z0-9]([A-Za-z0-9._-]*[A-Za-z0-9])?)")
def get_pyproject_package_names(pyproject_path: str, group: str) -> set[str]:
"""Return normalized names of packages explicitly listed in pyproject.toml."""
try:
import tomllib
except ModuleNotFoundError:
try:
import tomli as tomllib
except ModuleNotFoundError:
return set()
with open(pyproject_path, "rb") as f:
data = tomllib.load(f)
project = data.get("project", {})
reqs = list(project.get("dependencies", []))
reqs.extend(project.get("optional-dependencies", {}).get(group, []))
names = set()
for req in reqs:
match = _REQ_NAME_RE.match(req.strip())
if match:
names.add(_normalize_package_name(match.group(1)))
return names
def get_container_provided_packages() -> list[str]:
"""Discover torch/NVIDIA packages already installed in the container."""
from importlib.metadata import distributions
prefixes = ("torch", "torchvision", "triton", "nvidia-")
return sorted({
dist.metadata["Name"].lower()
for dist in distributions()
if dist.metadata["Name"].lower().startswith(prefixes)
})
def print_header():
"""Print setup header."""
log("=" * 60)
log(" TensorRT Diffusion - Dependency Setup")
log("=" * 60)
def check_uv_installed() -> tuple[bool, bool]:
"""
Check if uv is installed and accessible.
Returns:
tuple: (is_installed, needs_path_setup)
- is_installed: True if uv is available
- needs_path_setup: True if user needs to add ~/.local/bin to PATH permanently
"""
# First check if uv is in PATH
try:
subprocess.run(
["uv", "--version"],
capture_output=True,
check=True,
text=True
)
return True, False
except (subprocess.CalledProcessError, FileNotFoundError):
# If not in PATH, check the default installation location
local_bin = os.path.expanduser("~/.local/bin")
uv_path = os.path.join(local_bin, "uv")
if os.path.exists(uv_path) and os.access(uv_path, os.X_OK):
# uv exists but not in PATH - add it temporarily for this script
if local_bin not in os.environ["PATH"]:
os.environ["PATH"] = f"{local_bin}:{os.environ['PATH']}"
return True, True # Needs permanent PATH setup
return False, False
def upgrade_tensorrt(pip_spec: str) -> bool:
"""Upgrade/install TensorRT using pip.
Args:
pip_spec: PIP specifier for TensorRT, e.g. "tensorrt-cu12" or "tensorrt-cu12==10.16.*"
Returns:
True on success, False otherwise
"""
print("Installing TensorRT: {}".format(pip_spec))
try:
subprocess.run(
[sys.executable, "-m", "pip", "install", "--upgrade", "--pre", pip_spec],
check=True,
capture_output=_quiet,
)
# Print installed version for confirmation
try:
import importlib
trt = importlib.import_module("tensorrt")
log(f" Installed TensorRT version: {getattr(trt, '__version__', 'unknown')}")
except Exception:
pass
return True
except subprocess.CalledProcessError as e:
print(f"Failed to upgrade/install TensorRT: {e}")
return False
def install_libgl1() -> bool:
"""Install libgl1 system package when possible (Debian/Ubuntu).
Returns:
True if successfully installed or already present, False otherwise.
"""
print("Checking for libgl1 (system dependency)...")
apt_get = shutil.which("apt-get")
if not apt_get:
log(" Skipping: apt-get not found. Please install 'libgl1' via your OS package manager.")
return False
sudo = shutil.which("sudo")
use_sudo = False
try:
use_sudo = (hasattr(os, "geteuid") and os.geteuid() != 0 and sudo is not None)
except Exception:
use_sudo = sudo is not None
try:
# Update package list first
cmd_update = ([sudo] if use_sudo else []) + [apt_get, "update"]
subprocess.run(cmd_update, check=True, capture_output=_quiet)
cmd_install = ([sudo] if use_sudo else []) + [apt_get, "install", "-y", "libgl1"]
subprocess.run(cmd_install, check=True, capture_output=_quiet)
log(" libgl1 installed (or already up to date)")
return True
except subprocess.CalledProcessError as e:
print(f" Warning: Failed to install libgl1: {e}")
print(" Please install 'libgl1' manually via your OS package manager.")
return False
def install_uv():
"""Install uv package manager."""
log("Installing uv...")
try:
# Download and run uv installer
curl_cmd = [
"curl", "-LsSf",
"https://astral.sh/uv/install.sh"
]
curl_process = subprocess.Popen(
curl_cmd,
stdout=subprocess.PIPE,
stderr=subprocess.PIPE
)
sh_process = subprocess.Popen(
["sh"],
stdin=curl_process.stdout,
stdout=subprocess.PIPE,
stderr=subprocess.PIPE
)
curl_process.stdout.close()
stdout, stderr = sh_process.communicate()
if sh_process.returncode != 0:
print(f"Failed to install uv: {stderr.decode()}")
sys.exit(1)
# Add to PATH for this session
# uv installer uses ~/.local/bin by default
local_bin = os.path.expanduser("~/.local/bin")
if local_bin not in os.environ["PATH"]:
os.environ["PATH"] = f"{local_bin}:{os.environ['PATH']}"
log("uv installed successfully")
except Exception as e:
print(f"Error installing uv: {e}")
print("Please install manually: curl -LsSf https://astral.sh/uv/install.sh | sh")
sys.exit(1)
def get_installed_groups(deps_root: str) -> set[str]:
"""
Get set of already installed dependency groups.
A group is considered installed only if both:
1. The group directory exists
2. The .install_complete marker file exists in that directory
Args:
deps_root: Root directory where dependencies are installed
Returns:
Set of installed group names
"""
installed = set()
deps_path = Path(deps_root)
if not deps_path.exists():
return installed
for group in VALID_GROUPS:
group_dir = deps_path / group
marker_file = group_dir / INSTALL_COMPLETE_MARKER
# Only consider it installed if marker file exists
if group_dir.exists() and marker_file.exists():
installed.add(group)
return installed
def install_group(group: str, deps_root: str, project_root: str) -> bool:
"""
Install dependencies for a specific group.
Args:
group: Dependency group name
deps_root: Root directory where dependencies will be installed
project_root: Project root directory (where pyproject.toml is)
Returns:
True if installation succeeded, False otherwise
"""
description = GROUP_DESCRIPTIONS.get(group, group)
install_path = os.path.join(deps_root, group)
print(f"Installing {description}...")
log(f" Location: {install_path}")
try:
# Determine Python version and site-packages path
python_version = f"{sys.version_info.major}.{sys.version_info.minor}"
site_packages = Path(install_path) / "lib" / f"python{python_version}" / "site-packages"
# Create site-packages directory if it doesn't exist
site_packages.mkdir(parents=True, exist_ok=True)
# Build an overrides file so uv never installs packages that are
# already installed in the container. Exclude packages explicitly
# listed in pyproject.toml so they get installed at the pinned version.
container_pkgs = get_container_provided_packages()
declared_pkgs = get_pyproject_package_names(
str(Path(project_root) / "pyproject.toml"), group
)
overrides_content = "\n".join(
f'{pkg} ; python_version < "0"'
for pkg in container_pkgs
if _normalize_package_name(pkg) not in declared_pkgs
)
with tempfile.NamedTemporaryFile(
mode="w", suffix=".txt", prefix="uv_overrides_", delete=False
) as f:
f.write(overrides_content)
overrides_path = f.name
try:
cmd = [
"uv", "pip", "install",
"--python-preference", "only-system",
"--prefix", install_path,
"--overrides", overrides_path,
f".[{group}]"
]
subprocess.run(
cmd,
cwd=project_root,
text=True,
capture_output=_quiet,
check=True
)
finally:
os.unlink(overrides_path)
# Create marker file to indicate successful installation
marker_file = Path(install_path) / INSTALL_COMPLETE_MARKER
marker_file.write_text("Installation completed successfully\n")
print(f" {description} installed successfully")
return True
except subprocess.CalledProcessError as e:
print(f"Failed to install {description}")
print(f" Error: {e}")
# Clean up incomplete installation
if os.path.exists(install_path):
log(f" Cleaning up incomplete installation at {install_path}")
try:
shutil.rmtree(install_path)
except Exception as cleanup_error:
print(f" Warning: Failed to clean up: {cleanup_error}")
return False
def determine_groups_to_install(requested: str, already_installed: set[str]) -> list[str]:
"""
Determine which groups need to be installed.
Args:
requested: Requested group or "all"
already_installed: Set of already installed groups
Returns:
List of groups to install
"""
if requested == "all":
return sorted(VALID_GROUPS - already_installed)
elif requested in VALID_GROUPS:
return [] if requested in already_installed else [requested]
return []
def print_summary(installed_groups: set[str]):
"""
Print summary of installed groups.
Args:
installed_groups: Set of all installed groups
"""
log("=" * 60)
print("Setup complete!")
log("=" * 60)
if installed_groups:
log("Installed groups:")
for group in sorted(installed_groups):
description = GROUP_DESCRIPTIONS.get(group, group)
log(f" {description}")
log("Each demo script automatically uses the correct dependencies.")
else:
log("No groups installed.")
log()
def parse_args():
"""Parse command line arguments."""
parser = argparse.ArgumentParser(
description="Setup TensorRT Diffusion dependencies",
formatter_class=argparse.RawDescriptionHelpFormatter,
epilog="""
Examples:
python setup.py # Install all groups
python setup.py sd # Install SD family only
python setup.py flux # Install Flux only
python setup.py cosmos # Install Cosmos only
Groups:
sd - SD 1.4, SDXL, SD3, SD3.5, SVD, Stable Cascade (Stability AI)
flux - Flux models (Black Forest Labs)
cosmos - Cosmos models (NVIDIA)
all - All groups (default)
"""
)
parser.add_argument(
"group",
nargs="?",
default="all",
choices=list(VALID_GROUPS) + ["all"],
help="Dependency group to install (default: all)"
)
parser.add_argument(
"--deps-root",
default=DEFAULT_DEPS_ROOT,
help=f"Root directory for dependencies (default: {DEFAULT_DEPS_ROOT})"
)
parser.add_argument(
"--force",
action="store_true",
help="Force reinstallation even if already installed"
)
parser.add_argument("--skip-tensorrt", action="store_true", help="Skip TensorRT upgrade/installation")
parser.add_argument("-q", "--quiet", action="store_true", help="Suppress informational output (errors are always shown)")
return parser.parse_args()
def main():
"""Main setup function."""
args = parse_args()
global _quiet
_quiet = args.quiet
print_header()
# Check/install uv
uv_installed, needs_path_setup = check_uv_installed()
if not uv_installed:
install_uv()
needs_path_setup = True # Fresh install always needs PATH setup
else:
log("uv is already installed")
if needs_path_setup:
log(" (temporarily added ~/.local/bin to PATH for this session)")
# Get project root (where pyproject.toml is)
project_root = str(Path(__file__).parent.absolute())
pyproject_path = Path(project_root) / "pyproject.toml"
if not pyproject_path.exists():
print(f"Error: pyproject.toml not found in {project_root}")
print(" Make sure you're running this script from the project root.")
sys.exit(1)
# Check what's already installed
installed_groups = get_installed_groups(args.deps_root)
# Install system dependency (best-effort) and ensure TensorRT is present
install_libgl1()
if not args.skip_tensorrt:
upgrade_tensorrt("tensorrt-cu12")
else:
log("Skipping TensorRT upgrade (--skip-tensorrt specified)")
if installed_groups and not args.force:
log(f"Already installed: {', '.join(sorted(installed_groups))}")
# Determine what to install
if args.force and args.group == "all":
groups_to_install = sorted(VALID_GROUPS)
elif args.force and args.group in VALID_GROUPS:
groups_to_install = [args.group]
else:
groups_to_install = determine_groups_to_install(args.group, installed_groups)
# Early exit if nothing to do
if not groups_to_install:
if args.group == "all":
print("All requested dependencies are already installed!")
else:
description = GROUP_DESCRIPTIONS.get(args.group, args.group)
print(f"{description} is already installed!")
log(f" To reinstall, use --force flag or remove {args.deps_root}/{args.group}")
print_summary(installed_groups)
sys.exit(0)
# Install groups
print(f"Installing {len(groups_to_install)} group(s): {', '.join(groups_to_install)}")
success_count = 0
for group in groups_to_install:
if install_group(group, args.deps_root, project_root):
installed_groups.add(group)
success_count += 1
# Print summary
print_summary(installed_groups)
# Print PATH setup instructions if needed
if needs_path_setup:
print()
print("=" * 60)
print(" Action Required: Update PATH")
print("=" * 60)
print("uv was installed to ~/.local/bin but is not in your PATH.")
print()
print("To make uv available in future sessions, run:")
print(" source ~/.local/bin/env")
print()
print("Or restart your shell.")
print("=" * 60)
# Exit with error if any installations failed
if success_count < len(groups_to_install):
print(f"Warning: {len(groups_to_install) - success_count} group(s) failed to install")
sys.exit(1)
if __name__ == "__main__":
main()
demo/Diffusion/setup.sh
+21
View File
@@ -0,0 +1,21 @@
#!/bin/bash
# Minimal and backward-compatible: default to modern requirements, allow override via env.
REQ_FILE="${REQUIREMENTS_FILE:-requirements.txt}"
echo "[I] Using requirements file: ${REQ_FILE} (override with REQUIREMENTS_FILE=<file>)"
# Upgrade pip and install TensorRT
python3 -m pip install --upgrade pip
pip3 install --pre tensorrt-cu12
# Install the required packages
PIP_CONSTRAINT= pip3 install -r "${REQ_FILE}"
# Install libgl1
# Check if apt-get is available
if command -v apt-get >/dev/null 2>&1; then
sudo apt-get install -y libgl1
else
echo "Warning: apt-get not found. Please install libgl1 manually for your system."
fi
demo/Diffusion/tests/requirements.txt
+1
View File
@@ -0,0 +1 @@
pytest
demo/Diffusion/tests/test_utils_modelopt.py
+62
View File
@@ -0,0 +1,62 @@
#
# SPDX-FileCopyrightText: Copyright (c) 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 numpy as np
import onnx_graphsurgeon as gs
from demo_diffusion import utils_modelopt
def create_graph_with_fp16_resize() -> gs.Graph:
"""Return a gs.Graph with a single Resize node with FP16 input."""
return gs.Graph(
nodes=[
gs.Node(
op="Resize",
name="my_resize_node",
inputs=[
gs.Variable(name="X", dtype=np.float16),
gs.Variable(name="roi", dtype=np.float16),
gs.Variable(name="scales", dtype=np.float16),
gs.Variable(name="sizes", dtype=np.int64),
],
outputs=[gs.Variable(name="Y", dtype=np.float16)],
)
]
)
def test_should_cast_resize_to_fp32() -> None:
"""Test that `cast_resize_to_fp32` correctly casts all Resize nodes to FP32."""
# Precondition.
graph = create_graph_with_fp16_resize()
# Under test.
utils_modelopt.cast_resize_io(graph)
# Postcondition.
has_resize = False
for node in graph.nodes:
if node.op == "Resize":
has_resize = True
x, roi, scales, sizes = node.inputs
assert x.dtype == np.float32
assert roi.dtype == np.float32
assert scales.dtype == np.float32
assert sizes.dtype == np.int64 # "sizes" is the exception input that cannot be cast to FP32.
assert has_resize
demo/EfficientDet/notebooks/EfficientDet-TensorRT8.ipynb
+665
View File
File diff suppressed because one or more lines are too long
demo/EfficientDet/notebooks/README.md
+8
View File
@@ -0,0 +1,8 @@
# TensorRT demo with EfficientDet
To run the demo Jupyter notebook in this folder, follow the instructions in the TRT setup guide to generate the TensorRT-OSS docker container, and build TensorRT-OSS.
As specified in the guide for setting up the build environment, specify port number using `--jupyter <port>` to launch Jupyter lab interface with the container.
EfficientDet-TensorRT8.ipynb: Step by step walkthrough for building EfficientDet TensorRT engine from pre-trained Tensorflow checkpoint.