This demo showcases the Flux image generation model accelerated using Torch-TensorRT, with support for different precision modes (FP8, INT8, FP16) and dynamic shapes.

1. Install the required dependencies:

pip install gradio==5.29.0 nvidia-modelopt==0.27.1 diffusers==0.33.1 accelerate==1.3.0

The demo can be run with different configurations:

python flux_demo.py

• FP4 mode:

python flux_demo.py --dtype fp4

• FP8 mode:

python flux_demo.py --dtype fp8

• INT8 mode:

python flux_demo.py --dtype int8

• FP16 mode (default):

python flux_demo.py --dtype fp16

• Enable dynamic shapes (allows variable batch sizes):

python flux_demo.py --dynamic_shapes

• Low VRAM mode (for GPUs with ≤32GB VRAM):

python flux_demo.py --low_vram_mode

You can combine these options as needed. For example:

python flux_demo.py --dtype fp8 --dynamic_shapes --low_vram_mode