search

NVIDIA / TensorRT

NVIDIA® TensorRT™ is an SDK for high-performance deep learning inference on NVIDIA GPUs. This repository contains the open source components of TensorRT.
https://developer.nvidia.com/tensorrt
Apache License 2.0
10.75k stars 2.13k forks source link

Add Tensorrt conversion and inference for StableDiffusion Image2Image pipeline in demo #2688

Closed vanewu closed 1 year ago

vanewu commented 1 year ago

Description

Currently in this demo-diffusion.py , the VAE model conversion of StableDiffusion is only supported for VAE Decoder. Therefore, only the Text2Image pipeline is supported in this program. In daily use, the Image2Image pipeline is also often needed. I made some modifications based on the demo-diffusion.py , added the Image2Image part, and made some modifications to the model conversion and inference program. At present, the conversion has been successfully completed and a goodresult image has been obtained.

Can I submit this modification to the current demo?

Some modifications

  1. The VAE class in models.py is divided into two classes: VAEEncoder and VAEDecoder
  2. Add add_noise method to scheduler in utilities.py
  3. Add an optional memory_pool_limits parameter to the Engine build method in utilities.py
  4. Make some changes in demo-diffusion.py to support image2image inference, such as adjusting timesteps according to strength, adding init_image, strength and other parameters to infer

Rest of the items:

  1. Additionally, I found while using this script that it might be convenient to have an optional transformation of the model. For example, many times we may just want to convert the Unet model with different parameters, but keep the VAE or CLIP unchanged. It seems more convenient to provide an optional name parameter for select conversion.

  2. When force_export is true, maybe you want to re-export the onnx model, but the existence of onnx_opt_path will make it impossible to re-export and optimize onnx. Does this logic need to be modified, when force_export is true, should the onnx model be forcibly exported even if the file exists.

Relevant Files

demo/diffusion

zerollzeng commented 1 year ago

@rajeevsrao ^ ^

venetanji commented 1 year ago

+1 for this. Could you share the code or make a pr? I've also successfully compiled a vaeencode engine, but the encoding itself had issues. It was messing up the resolution somehow, overlapping different areas of the original image. I've also separated the onnx export and engine compile script. You can check my code here: https://github.com/venetanji/videosd

rajeevsrao commented 1 year ago

@vanewu @venetanji the new demoDiffusion code in release/8.6 branch supports img2img and inpainting as well.