gpu-trt-infer-demo.py

# Copyright 2022 The OFA-Sys Team.
# All rights reserved.
# This source code is licensed under the Apache 2.0 license
# found in the LICENSE file in the root directory.
# SPDX-FileCopyrightText: Copyright (c) 1993-2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#

import argparse
from cuda import cudart
from models import CLIP, UNet, VAE
import numpy as np
import nvtx
import os
import onnx
from polygraphy import cuda
import time
import torch
from transformers import CLIPTokenizer
import tensorrt as trt
from utilities import Engine, DPMScheduler, LMSDiscreteScheduler, save_image, TRT_LOGGER


def parseArgs():
    parser = argparse.ArgumentParser(
        description="Options for Stable Diffusion Demo")
    # Stable Diffusion configuration
    parser.add_argument('prompt',
                        nargs='*',
                        help="Text prompt(s) to guide image generation")
    parser.add_argument(
        '--height',
        type=int,
        default=512,
        help="Height of image to generate (must be multiple of 8)")
    parser.add_argument(
        '--width',
        type=int,
        default=512,
        help="Height of image to generate (must be multiple of 8)")
    parser.add_argument('--num-images',
                        type=int,
                        default=1,
                        help="Number of images to generate per prompt")
    parser.add_argument('--steps',
                        type=int,
                        default=50,
                        help="Number of inference steps")
    parser.add_argument('--denoise-prec',
                        type=str,
                        default='fp16',
                        choices=['fp32', 'fp16'],
                        help="UNet model precision")
    parser.add_argument(
        '--negative-prompt',
        nargs='*',
        default=[''],
        help="The negative prompt(s) to guide the image generation.")
    parser.add_argument(
        '--repeat-prompt',
        type=int,
        default=1,
        choices=[1, 2, 4, 8, 16],
        help="Number of times to repeat the prompt (batch size multiplier)")
    parser.add_argument('--scheduler',
                        type=str,
                        default="LMSD",
                        choices=["LMSD", "DPM"],
                        help="Scheduler for diffusion process")

    # ONNX export
    parser.add_argument(
        '--onnx-opset',
        type=int,
        default=16,
        choices=range(7, 18),
        help="Select ONNX opset version to target for exported models")
    parser.add_argument('--onnx-dir',
                        default='onnx',
                        help="Output directory for ONNX export")
    parser.add_argument('--force-onnx-export',
                        action='store_true',
                        help="Force ONNX export of CLIP, UNET, and VAE models")
    parser.add_argument(
        '--force-onnx-optimize',
        action='store_true',
        help="Force ONNX optimizations for CLIP, UNET, and VAE models")
    parser.add_argument('--force-engine-build',
                        action='store_true',
                        help="Force rebuilding the TensorRT engine")
    parser.add_argument(
        '--force-static-batch',
        action='store_true',
        help="Force building TensorRT engines with fixed batch size.")
    parser.add_argument(
        '--minimal-optimization',
        action='store_true',
        help="Limited optimizations to only const folding and shape inference."
    )
    parser.add_argument('--enable-preview-features',
                        action='store_true',
                        help="Enable TensorRT preview features.")

    # TensorRT inference
    parser.add_argument('--engine-dir',
                        default='engine',
                        help="Output directory for TensorRT engines")
    parser.add_argument(
        '--num-warmup-runs',
        type=int,
        default=5,
        help="Number of warmup runs before benchmarking performance")
    parser.add_argument('--profile',
                        action='store_true',
                        help="Enable performance profiling")
    parser.add_argument(
        '--seed',
        type=int,
        default=None,
        help="Seed for random generator to get consistent results")

    parser.add_argument('--output-dir',
                        default='output',
                        help="Output directory for logs and image artifacts")
    parser.add_argument(
        '--hf-token',
        type=str,
        help="HuggingFace API token to use for downloading checkpoints")
    parser.add_argument('-v',
                        '--verbose',
                        action='store_true',
                        help="Show verbose output")
    return parser.parse_args()


class DemoDiffusion:
    """
    Application showcasing the acceleration of [Stable Diffusion v1.4](https://huggingface.co/CompVis/stable-diffusion-v1-4) pipeline using NVidia TensorRT w/ Plugins.
    """
    def __init__(
        self,
        image_height,
        image_width,
        denoising_steps,
        scheduler="LMSD",
        denoising_fp16=True,
        guidance_scale=7.5,
        device='cuda',
        output_dir='.',
        hf_token=None,
        verbose=False,
        profile=False,
    ):
        """
        Initializes the Diffusion pipeline.

        Args:
            image_height (int):
                Height (in pixels) of the image to be generated. Should be a multiple of 8.
            image_width (int):
                Width (in pixels) of the image to be generated. Should be a multiple of 8.
            denoising_steps (int):
                The number of denoising steps.
                More denoising steps usually lead to a higher quality image at the expense of slower inference.
            denoising_fp16 (bool):
                Run the denoising loop (UNet) in fp16 precision.
                When enabled image quality will be lower but generally results in higher throughput.
            guidance_scale (float):
                Guidance scale is enabled by setting as > 1.
                Higher guidance scale encourages to generate images that are closely linked to the text prompt, usually at the expense of lower image quality.
            device (str):
                PyTorch device to run inference. Default: 'cuda'
            output_dir (str):
                Output directory for log files and image artifacts
            hf_token (str):
                HuggingFace User Access Token to use for downloading Stable Diffusion model checkpoints.
            verbose (bool):
                Enable verbose logging.
            profile (bool):
                Insert NVTX profiling markers.
        """

        if image_height % 8 != 0 or image_width % 8 != 0:
            raise ValueError(
                f"Image height and width have to be divisible by 8 but specified as: {image_height} and {image_width}."
            )
        # Spatial dimensions of latent tensor
        self.latent_height = image_height // 8
        self.latent_width = image_width // 8

        # Only supports single image per prompt.
        self.num_images = 1

        self.denoising_steps = denoising_steps
        self.denoising_fp16 = denoising_fp16
        assert guidance_scale > 1.0
        self.guidance_scale = guidance_scale

        self.output_dir = output_dir
        self.hf_token = hf_token
        self.device = device
        self.verbose = verbose
        self.profile = profile

        # A scheduler to be used in combination with unet to denoise the encoded image latens.
        # This demo uses an adaptation of LMSDiscreteScheduler or DPMScheduler:
        sched_opts = {
            'num_train_timesteps': 1000,
            'beta_start': 0.00085,
            'beta_end': 0.012
        }
        if scheduler == "DPM":
            self.scheduler = DPMScheduler(device=self.device, **sched_opts)
        elif scheduler == "LMSD":
            self.scheduler = LMSDiscreteScheduler(device=self.device,
                                                  **sched_opts)
        else:
            raise ValueError(f"Scheduler should be either DPM or LMSD")

        self.tokenizer = None

        self.unet_model_key = 'unet_fp16' if denoising_fp16 else 'unet'
        self.models = {
            'clip':
            CLIP(hf_token=hf_token,
                 image_width=image_width,
                 image_height=image_height,
                 device=device,
                 verbose=verbose),
            self.unet_model_key:
            UNet(hf_token=hf_token,
                 image_width=image_width,
                 image_height=image_height,
                 fp16=denoising_fp16,
                 device=device,
                 verbose=verbose),
            'vae':
            VAE(hf_token=hf_token,
                image_width=image_width,
                image_height=image_height,
                device=device,
                verbose=verbose)
        }

        self.engine = {}
        self.stream = cuda.Stream()

    def teardown(self):
        for engine in self.engine.values():
            del engine
        self.stream.free()
        del self.stream

    def getModelPath(self, name, onnx_dir, opt=True):
        return os.path.join(onnx_dir, name + ('.opt' if opt else '') + '.onnx')

    def loadEngines(
        self,
        engine_dir,
        onnx_dir,
        onnx_opset,
        opt_batch_size,
        force_export=False,
        force_optimize=False,
        force_build=False,
        minimal_optimization=False,
        static_batch=False,
        enable_preview=False,
    ):
        """
        Build and load engines for TensorRT accelerated inference.
        Export ONNX models first, if applicable.

        Args:
            engine_dir (str):
                Directory to write the TensorRT engines.
            onnx_dir (str):
                Directory to write the ONNX models.
            onnx_opset (int):
                ONNX opset version to export the models.
            opt_batch_size (int):
                Batch size to optimize for during engine building.
            force_export (bool):
                Force re-exporting the ONNX models.
            force_optimize (bool):
                Force re-optimizing the ONNX models.
            force_build (bool):
                Force re-building the TensorRT engine.
            minimal_optimization (bool):
                Apply minimal optimizations during build (no plugins).
            static_batch (bool):
                Build engine only for specified opt_batch_size.
            enable_preview (bool):
                Enable TensorRT preview features.
        """
        def exportOnnx(model_name, obj):
            onnx_opt_path = self.getModelPath(model_name, onnx_dir)
            if force_optimize or not os.path.exists(onnx_opt_path):
                onnx_path = self.getModelPath(model_name, onnx_dir, opt=False)
                if force_export or not os.path.exists(onnx_path):
                    print(f"Exporting model: {onnx_path}")
                    model = obj.get_model()
                    with torch.inference_mode(), torch.autocast("cuda"):
                        inputs = obj.get_sample_input(
                            batch_size=opt_batch_size)
                        torch.onnx.export(
                            model,
                            inputs,
                            onnx_path,
                            export_params=True,
                            opset_version=onnx_opset,
                            do_constant_folding=True,
                            input_names=obj.get_input_names(),
                            output_names=obj.get_output_names(),
                            dynamic_axes=obj.get_dynamic_axes(),
                        )
                else:
                    print(f"Found cached model: {onnx_path}")

                print(f"Generating optimizing model: {onnx_opt_path}")
                onnx_opt_graph = obj.optimize(
                    onnx.load(onnx_path),
                    minimal_optimization=minimal_optimization)
                onnx.save(onnx_opt_graph, onnx_opt_path)
            else:
                print(f"Found cached optimized model: {onnx_opt_path} ")

        # Build engines
        for model_name, obj in self.models.items():
            engine = Engine(model_name, engine_dir)
            if force_build or not os.path.exists(engine.engine_path):
                onnx_path = self.getModelPath(model_name, onnx_dir)
                if not os.path.exists(onnx_path):
                    exportOnnx(model_name, obj)
                engine.build(onnx_path, fp16=True, \
                    input_profile=obj.get_input_profile(batch_size=opt_batch_size, static_batch=static_batch), \
                    enable_preview=enable_preview)

        # Load engines
        for model_name, obj in self.models.items():
            engine = Engine(model_name, engine_dir)
            engine.activate()
            self.engine[model_name] = engine

    def loadModules(self, ):
        self.tokenizer = CLIPTokenizer.from_pretrained(
            "openai/clip-vit-large-patch14")
        self.scheduler.set_timesteps(self.denoising_steps)
        # Pre-compute latent input scales and linear multistep coefficients
        self.scheduler.configure()

    def runEngine(self, model_name, feed_dict):
        engine = self.engine[model_name]
        return engine.infer(feed_dict, self.stream)

    def infer(
        self,
        prompt,
        negative_prompt,
        warmup=False,
        verbose=False,
    ):
        """
        Run the diffusion pipeline.

        Args:
            prompt (str):
                The text prompt to guide image generation.
            negative_prompt (str):
                The prompt not to guide the image generation.
            warmup (bool):
                Indicate if this is a warmup run.
            verbose (bool):
                Enable verbose logging.
        """
        # Process inputs
        batch_size = len(prompt)
        assert len(prompt) == len(negative_prompt)

        # Create profiling events
        events = {}
        for stage in ['clip', 'denoise', 'vae']:
            for marker in ['start', 'stop']:
                events[stage + '-' + marker] = cudart.cudaEventCreate()[1]

        # Allocate buffers for TensorRT engine bindings
        for model_name, obj in self.models.items():
            self.engine[model_name].allocate_buffers(
                shape_dict=obj.get_shape_dict(batch_size=batch_size),
                device=self.device)

        generator = None
        if args.seed is not None:
            generator = torch.Generator(device="cuda").manual_seed(args.seed)

        # Run Stable Diffusion pipeline
        with torch.inference_mode(), torch.autocast("cuda"), trt.Runtime(
                TRT_LOGGER) as runtime:
            # latents need to be generated on the target device
            unet_channels = 4  # unet.in_channels
            latents_shape = (batch_size * self.num_images, unet_channels,
                             self.latent_height, self.latent_width)
            latents_dtype = torch.float32  # text_embeddings.dtype
            latents = torch.randn(latents_shape,
                                  device=self.device,
                                  dtype=latents_dtype,
                                  generator=generator)

            # Scale the initial noise by the standard deviation required by the scheduler
            latents = latents * self.scheduler.init_noise_sigma

            torch.cuda.synchronize()
            e2e_tic = time.perf_counter()

            if self.profile:
                nvtx_clip = nvtx.start_range(message='clip', color='green')
            cudart.cudaEventRecord(events['clip-start'], 0)
            # Tokenize input
            text_input_ids = self.tokenizer(
                prompt,
                padding="max_length",
                max_length=self.tokenizer.model_max_length,
                return_tensors="pt",
            ).input_ids.type(torch.int32).to(self.device)

            # CLIP text encoder
            text_input_ids_inp = cuda.DeviceView(ptr=text_input_ids.data_ptr(),
                                                 shape=text_input_ids.shape,
                                                 dtype=np.int32)
            text_embeddings = self.runEngine(
                'clip', {"input_ids": text_input_ids_inp})['text_embeddings']

            # Duplicate text embeddings for each generation per prompt
            bs_embed, seq_len, _ = text_embeddings.shape
            text_embeddings = text_embeddings.repeat(1, self.num_images, 1)
            text_embeddings = text_embeddings.view(bs_embed * self.num_images,
                                                   seq_len, -1)

            max_length = text_input_ids.shape[-1]
            uncond_input_ids = self.tokenizer(
                negative_prompt,
                padding="max_length",
                max_length=max_length,
                truncation=True,
                return_tensors="pt",
            ).input_ids.type(torch.int32).to(self.device)
            uncond_input_ids_inp = cuda.DeviceView(
                ptr=uncond_input_ids.data_ptr(),
                shape=uncond_input_ids.shape,
                dtype=np.int32)
            uncond_embeddings = self.runEngine(
                'clip', {"input_ids": uncond_input_ids_inp})['text_embeddings']

            # Duplicate unconditional embeddings for each generation per prompt
            seq_len = uncond_embeddings.shape[1]
            uncond_embeddings = uncond_embeddings.repeat(1, self.num_images, 1)
            uncond_embeddings = uncond_embeddings.view(
                batch_size * self.num_images, seq_len, -1)

            # Concatenate the unconditional and text embeddings into a single batch to avoid doing two forward passes for classifier free guidance
            text_embeddings = torch.cat([uncond_embeddings, text_embeddings])

            if self.denoising_fp16:
                text_embeddings = text_embeddings.to(dtype=torch.float16)

            cudart.cudaEventRecord(events['clip-stop'], 0)
            if self.profile:
                nvtx.end_range(nvtx_clip)

            cudart.cudaEventRecord(events['denoise-start'], 0)
            for step_index, timestep in enumerate(self.scheduler.timesteps):
                if self.profile:
                    nvtx_latent_scale = nvtx.start_range(
                        message='latent_scale', color='pink')
                # expand the latents if we are doing classifier free guidance
                latent_model_input = torch.cat([latents] * 2)
                # LMSDiscreteScheduler.scale_model_input()
                latent_model_input = self.scheduler.scale_model_input(
                    latent_model_input, step_index)
                if self.profile:
                    nvtx.end_range(nvtx_latent_scale)

                # predict the noise residual
                if self.profile:
                    nvtx_unet = nvtx.start_range(message='unet', color='blue')
                dtype = np.float16 if self.denoising_fp16 else np.float32
                if timestep.dtype != torch.float32:
                    timestep_float = timestep.float()
                else:
                    timestep_float = timestep
                sample_inp = cuda.DeviceView(ptr=latent_model_input.data_ptr(),
                                             shape=latent_model_input.shape,
                                             dtype=np.float32)
                timestep_inp = cuda.DeviceView(ptr=timestep_float.data_ptr(),
                                               shape=timestep_float.shape,
                                               dtype=np.float32)
                embeddings_inp = cuda.DeviceView(
                    ptr=text_embeddings.data_ptr(),
                    shape=text_embeddings.shape,
                    dtype=dtype)
                noise_pred = self.runEngine(
                    self.unet_model_key, {
                        "sample": sample_inp,
                        "timestep": timestep_inp,
                        "encoder_hidden_states": embeddings_inp
                    })['latent']
                if self.profile:
                    nvtx.end_range(nvtx_unet)

                if self.profile:
                    nvtx_latent_step = nvtx.start_range(message='latent_step',
                                                        color='pink')
                # Perform guidance
                noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
                noise_pred = noise_pred_uncond + self.guidance_scale * (
                    noise_pred_text - noise_pred_uncond)

                latents = self.scheduler.step(noise_pred, latents, step_index,
                                              timestep)

                if self.profile:
                    nvtx.end_range(nvtx_latent_step)

            latents = 1. / 0.18215 * latents
            cudart.cudaEventRecord(events['denoise-stop'], 0)

            if self.profile:
                nvtx_vae = nvtx.start_range(message='vae', color='red')
            cudart.cudaEventRecord(events['vae-start'], 0)
            sample_inp = cuda.DeviceView(ptr=latents.data_ptr(),
                                         shape=latents.shape,
                                         dtype=np.float32)
            images = self.runEngine('vae', {"latent": sample_inp})['images']
            cudart.cudaEventRecord(events['vae-stop'], 0)
            if self.profile:
                nvtx.end_range(nvtx_vae)

            torch.cuda.synchronize()
            e2e_toc = time.perf_counter()
            if not warmup:
                print('|------------|--------------|')
                print('| {:^10} | {:^12} |'.format('Module', 'Latency'))
                print('|------------|--------------|')
                print('| {:^10} | {:>9.2f} ms |'.format(
                    'CLIP',
                    cudart.cudaEventElapsedTime(events['clip-start'],
                                                events['clip-stop'])[1]))
                print('| {:^10} | {:>9.2f} ms |'.format(
                    'UNet x ' + str(self.denoising_steps),
                    cudart.cudaEventElapsedTime(events['denoise-start'],
                                                events['denoise-stop'])[1]))
                print('| {:^10} | {:>9.2f} ms |'.format(
                    'VAE',
                    cudart.cudaEventElapsedTime(events['vae-start'],
                                                events['vae-stop'])[1]))
                print('|------------|--------------|')
                print('| {:^10} | {:>9.2f} ms |'.format(
                    'Pipeline', (e2e_toc - e2e_tic) * 1000.))
                print('|------------|--------------|')

                # Save image
                image_name_prefix = 'sd-' + (
                    'fp16' if self.denoising_fp16 else 'fp32') + ''.join(
                        set([
                            '-' + prompt[i].replace(' ', '_')[:10]
                            for i in range(batch_size)
                        ])) + '-'
                save_image(images, self.output_dir, image_name_prefix)


if __name__ == "__main__":

    print("[I] Initializing StableDiffusion demo with TensorRT Plugins")
    args = parseArgs()

    # Process prompt
    if not isinstance(args.prompt, list):
        raise ValueError(
            f"`prompt` must be of type `str` or `str` list, but is {type(prompt)}"
        )
    prompt = args.prompt * args.repeat_prompt

    if not isinstance(args.negative_prompt, list):
        raise ValueError(
            f"`--negative-prompt` must be of type `str` or `str` list, but is {type(args.negative_prompt)}"
        )
    if len(args.negative_prompt) == 1:
        negative_prompt = args.negative_prompt * len(prompt)
    else:
        negative_prompt = args.negative_prompt

    # Register TensorRT plugins
    trt.init_libnvinfer_plugins(TRT_LOGGER, '')

    # Initialize demo
    demo = DemoDiffusion(image_height=args.height,
                         image_width=args.width,
                         denoising_steps=args.steps,
                         denoising_fp16=(args.denoise_prec == 'fp16'),
                         output_dir=args.output_dir,
                         scheduler=args.scheduler,
                         hf_token=args.hf_token,
                         verbose=args.verbose,
                         profile=args.profile)

    # Build/load TensorRT engines and torch models
    demo.loadEngines(args.engine_dir, args.onnx_dir, args.onnx_opset, opt_batch_size=len(prompt), \
        force_export=args.force_onnx_export, force_optimize=args.force_onnx_optimize, \
        force_build=args.force_engine_build, minimal_optimization=args.minimal_optimization, \
        static_batch=args.force_static_batch, enable_preview=args.enable_preview_features)
    demo.loadModules()

    print("[I] Warming up ..")
    for _ in range(args.num_warmup_runs):
        images = demo.infer(prompt,
                            negative_prompt,
                            warmup=True,
                            verbose=False)

    print("[I] Running StableDiffusion pipeline")
    if args.profile:
        cudart.cudaProfilerStart()
    images = demo.infer(prompt, negative_prompt, verbose=args.verbose)
    if args.profile:
        cudart.cudaProfilerStop()

    demo.teardown()