Add trt support for BF16

[andompesta]

This pull request aims to add support for TensortRT integration

Summary of Changes

1. CLI Support for TensorRT

   • Added a new variable in src/flux/cli.py to enable TensorRT inference.
   • environment variable TRT_ENGINE_DIR specifies the directory for storing TensorRT engines
   • environment variable ONNX_DIRspecifies the directory for ONNX model exports.
   • surrently supports the bf16 (fp16 and fp8 coming soon)
   • supports for model offloading can be added

2. Modifications for ONNX Export

   • Minor changes to src/flux/math and src/flux/modules/autoencoder.py to enable proper export in ONNX
   • Additional changes are required to address numerical stability issues within the Flux-Transformer model.

3. TensorRT Exporter

   • Added the flux/trt/exporter package, containing code to export PyTorch models to ONNX and build TensorRT engines.

4. TensorRT Engine Execution

   • flux/trt/engine package is responsible to execute inference using TRT

5. TensorRT Mixin Classes

   • Added the flux/trt/mixin package with mixin classes to share parameters between the model building and inference phases.

6. TensorRT Manager

   • Introduced flux/trt/trt_manager.py as the main TensorRT management class. It handles the conversion of PyTorch models to TensorRT engines and manages the TensorRT context for inference.