Model Export to TensorRT-LLM

For DreamBooth, the export script generates three different optimized inference models. The first model is the VAE Decoder, the second model is the UNet, and the third model is the CLIP Encoder.

  1. In the defaults section of conf/config.yaml, update the export field to point to the desired Stable Diffusion inference configuration file. For example, if you want to use the dreambooth/export_dreambooth.yaml configuration, change the export field to dreambooth/export_dreambooth.

    Copy
    Copied!
                

    defaults: - export: dreambooth/export_dreambooth ...

  2. In the stages field of conf/config.yaml, make sure the export stage is included. For example,

    Copy
    Copied!
                

    stages: - export ...

  3. Configure infer.num_images_per_prompt of the conf/export/dreambooth/export_dreambooth.yaml file to set the batch_size to use for the ONNX and NVIDIA TensorRT models.

Remarks:

  1. To load a pretrained checkpoint for inference, set the restore_from_path field in the model section to the path of the pretrained checkpoint in .nemo format in conf/export/dreambooth/export_dreambooth.yaml.

Previous Framework Inference
Next Model Deployment
© Copyright 2023-2024, NVIDIA. Last updated on May 17, 2024.