# 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`.

