home/ atoms/ tensorrt-acceleration-for-diffusion

Compiling UNet and VAE to TensorRT engines delivers the largest throughput gain in StreamDiffusion

StreamDiffusion supports three acceleration modes: none (baseline PyTorch), xformers (memory-efficient attention), and tensorrt (compiled engines). The TensorRT path exports the UNet, VAE encoder, and VAE decoder to ONNX, optimises each graph, and builds GPU-specific .engine files. At runtime the PyTorch UNet and VAE are replaced with UNet2DConditionModelEngine and AutoencoderKLEngine wrappers that call the compiled engines via CUDA streams. Engine compilation is done once per model+batch-size combination and cached to disk; subsequent runs load pre-compiled engines. On RTX 4090, TensorRT mode yields SD-Turbo t2i at 106 fps vs lower rates without it. The tradeoff: compilation takes minutes and engines are hardware-specific.

Examples

from streamdiffusion.acceleration.tensorrt import accelerate_with_tensorrt

stream = accelerate_with_tensorrt(
    stream, 'engines', max_batch_size=2
)
# or via wrapper:
streamwrapper = StreamDiffusionWrapper(
    ..., acceleration='tensorrt'
)

Assessment

Why must TensorRT engines be rebuilt when switching from img2img to txt2img mode? Identify what parameters determine the engine filename prefix in the wrapper’s create_prefix() function.

“It requires TensorRT extension and time to build the engine, but it will be faster than the above example.”
corpus · streamdiffusion-pipeline-for-real-time-interactive-image-gen · chunk 6