[feature] Support Flux TensorRT Pipeline #12218
Conversation
Co-authored-by: Steven Liu <[email protected]>
Co-authored-by: Steven Liu <[email protected]>
Co-authored-by: Steven Liu <[email protected]>
|
cc @sayakpaul |
examples/flux-tensorrt/README.md
Outdated
| engine_transformer_path = "path/to/transformer/engine_trt10.13.2.6.plan" | ||
| engine_vae_path = "path/to/vae/engine_trt10.13.2.6.plan" | ||
| engine_t5xxl_path = "path/to/t5/engine_trt10.13.2.6.plan" | ||
| engine_clip_path = "path/to/clip/engine_trt10.13.2.6.plan" |
There was a problem hiding this comment.
How to derive these? Can we host some of these files on the Hub and supplement in this project?
There was a problem hiding this comment.
These .plan files depend on the GPUs that we use. I only have them for H100 SXM, do you need these files to run?
examples/flux-tensorrt/README.md
Outdated
| We follow the official [NVIDIA/TensorRT](https://github.com/NVIDIA/TensorRT) repository to build TensorRT. | ||
|
|
||
| > **Note:** | ||
| > TensorRT was originally built with `diffusers==0.31.1`. |
There was a problem hiding this comment.
TensorRT's build shouldn't depend on the diffusers version no?
There was a problem hiding this comment.
Yeah, but in the TensorRT repo, it's used with this version. However, after building the TRT files, we can use my script for inference with the current diffusers version.
There was a problem hiding this comment.
maybe will have other ways for build or I will code a new one. This is just fast building.
examples/flux-tensorrt/README.md
Outdated
| > - one **venv** for building, and | ||
| > - another **venv** for inference. | ||
|
|
||
| (🔜 TODO: Build scripts for the latest `diffusers` will be added later.) |
There was a problem hiding this comment.
What is needed for this? Can we maybe comment on that so that other contributors could pick it up if interested?
examples/flux-tensorrt/README.md
Outdated
| pip install -r requirements.txt | ||
| ``` | ||
|
|
||
| ### ⚡ Fast Building with Static Shapes |
There was a problem hiding this comment.
If this can be run from the repository NVIDIA/TensorRT itself, then what is the purpose of this example?
There was a problem hiding this comment.
Just a guide for building .plan files, in case someone has no idea how to build them.
There was a problem hiding this comment.
Oh then I think it should live in the NVIDIA TensorRT repository as most of the examples seem to be taken from there.
There was a problem hiding this comment.
Okay @sayakpaul, I will remove it and maybe write code to build it in Diffusers in the near future.
Co-authored-by: Sayak Paul <[email protected]>
|
hey @sayakpaul, I added code to build TRT. You can take a look. I think everything is working well now. |
| You can convert all ONNX checkpoints to TensorRT engines with a single command: | ||
| ```bash | ||
| python convert_trt.py | ||
| ``` |
There was a problem hiding this comment.
Do we need onnx checkpoints for TensorRT?
There was a problem hiding this comment.
yeah, at REAME, I mentioned it in this line.
There was a problem hiding this comment.
It could be nice to have the possibility of taking a diffusers checkpoint and then convert it to TensorRT.
Do you think that is possible?
There was a problem hiding this comment.
Yeah, I did it, but in my own way. I think I’ll code this for Diffusers and create a new PR.
There was a problem hiding this comment.
That would be actually amazing!
Co-authored-by: Sayak Paul <[email protected]>
|
The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update. |
|
hey @sayakpaul , did this check have sth wrong? |
|
Not sure why did you close the PR. |
|
Sorry, I’m not sure why this check failed. I asked you about it earlier but didn’t get a reply. I think there’s something wrong with my PR, so I closed it |
|
@bot /style |
|
Style bot fixed some files and pushed the changes. |
|
Thanks @sayakpaul for reopening! I’ll create a new PR for ONNX export once this one is merged into main |
3 participants
What does this PR do?
This PR addresses issue #12202.
It introduces initial support for the
FluxPipelinewith TensorRT acceleration.For installation steps and usage examples, please check the updated README.md.
Below are sample results comparing PyTorch (before TRT) vs TensorRT (after TRT):
🐱 Example 1:
👧 Example 2: