nvidia--tensorrt/demo/Diffusion/demo_stable_cascade.py

#
# SPDX-FileCopyrightText: Copyright (c) 1993-2025 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.
#

# Configure dependencies before any external imports
from demo_diffusion import deps
deps.configure("sd")

import argparse
import os

import torch
from cuda.bindings import runtime as cudart

from demo_diffusion import dd_argparse
from demo_diffusion import pipeline as pipeline_module


def parseArgs():
    parser = argparse.ArgumentParser(description="Options for Stable Cascade Txt2Img Demo", conflict_handler='resolve')
    parser = dd_argparse.add_arguments(parser)
    parser.add_argument('--version', type=str, default="cascade", choices=["cascade"], help="Version of Stable Cascade")
    parser.add_argument('--height', type=int, default=1024, help="Height of image to generate (must be multiple of 8)")
    parser.add_argument('--width', type=int, default=1024, help="Width of image to generate (must be multiple of 8)")
    parser.add_argument('--lite', action='store_true', help="Use the Lite Version of the Stage B and Stage C models")
    parser.add_argument('--prior-guidance-scale', type=float, default=4.0, help="Value of classifier-free guidance scale for the prior")
    parser.add_argument('--decoder-guidance-scale', type=float, default=0.0, help="Value of classifier-free guidance scale for the decoder")
    parser.add_argument('--prior-denoising-steps', type=int, default=20, help="Number of denoising steps for the prior")
    parser.add_argument('--decoder-denoising-steps', type=int, default=10, help="Number of denoising steps for the decoder")
    return parser.parse_args()


class StableCascadeDemoPipeline(pipeline_module.StableCascadePipeline):
    def __init__(self, prior_denoising_steps, decoder_denoising_steps, prior_guidance_scale, decoder_guidance_scale, lite, **kwargs):
        self.nvtx_profile = kwargs['nvtx_profile']
        self.prior = pipeline_module.StableCascadePipeline(
            pipeline_type=pipeline_module.PIPELINE_TYPE.CASCADE_PRIOR,
            denoising_steps=prior_denoising_steps,
            guidance_scale=prior_guidance_scale,
            return_latents=True,
            lite=lite,
            **kwargs,
        )
        self.decoder = pipeline_module.StableCascadePipeline(
            pipeline_type=pipeline_module.PIPELINE_TYPE.CASCADE_DECODER,
            denoising_steps=decoder_denoising_steps,
            guidance_scale=decoder_guidance_scale,
            lite=lite,
            **kwargs,
        )

    def loadEngines(self, framework_model_dir, onnx_dir, engine_dir, **kwargs):
        prior_suffix = "prior_lite" if self.prior.lite else "prior"
        decoder_suffix = "decoder_lite" if self.decoder.lite else "decoder"
        self.prior.loadEngines(
            os.path.join(engine_dir, prior_suffix),
            framework_model_dir,
            os.path.join(onnx_dir, prior_suffix),
            **kwargs)
        self.decoder.loadEngines(
            os.path.join(engine_dir, decoder_suffix),
            framework_model_dir,
            os.path.join(onnx_dir, decoder_suffix),
            **kwargs)

    def activateEngines(self, shared_device_memory=None):
        self.prior.activateEngines(shared_device_memory)
        self.decoder.activateEngines(shared_device_memory)

    def loadResources(self, image_height, image_width, batch_size, seed):
        self.prior.loadResources(image_height, image_width, batch_size, seed)
        # Use a different seed for decoder
        self.decoder.loadResources(image_height, image_width, batch_size, ((seed+1) if seed is not None else None))

    def get_max_device_memory(self):
        max_device_memory = self.prior.calculateMaxDeviceMemory()
        max_device_memory = max(max_device_memory, self.decoder.calculateMaxDeviceMemory())
        return max_device_memory

    def run(self, prompt, negative_prompt, height, width, batch_size, batch_count, num_warmup_runs, use_cuda_graph):
        # Process prompt
        if not isinstance(prompt, list):
            raise ValueError(f"`prompt` must be of type `str` list, but is {type(prompt)}")
        prompt = prompt * batch_size

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

        num_warmup_runs = max(1, num_warmup_runs) if use_cuda_graph else num_warmup_runs
        if num_warmup_runs > 0:
            print("[I] Warming up ..")
            for _ in range(num_warmup_runs):
                latents, _ = self.prior.infer(prompt, negative_prompt, height, width, warmup=True)
                latents = latents.to(torch.float16) if self.decoder.fp16 else latents
                images, _ = self.decoder.infer(prompt, negative_prompt, height, width, image_embeddings=latents, warmup=True)

        for _ in range(batch_count):
            print("[I] Running Stable Cascade pipeline")
            if self.nvtx_profile:
                cudart.cudaProfilerStart()
            latents, time_prior = self.prior.infer(prompt, negative_prompt, height, width, warmup=False)
            latents = latents.to(torch.float16) if self.decoder.fp16 else latents
            images, time_decoder = self.decoder.infer(prompt, negative_prompt, height, width, image_embeddings=latents, warmup=False)

            if self.nvtx_profile:
                cudart.cudaProfilerStop()
            print('|-----------------|--------------|')
            print('| {:^15} | {:>9.2f} ms |'.format('e2e', time_prior + time_decoder))
            print('|-----------------|--------------|')

    def teardown(self):
        self.prior.teardown()
        self.decoder.teardown()


if __name__ == "__main__":
    print("[I] Initializing StableCascade txt2img demo using TensorRT")
    args = parseArgs()

    kwargs_init_pipeline, kwargs_load_engine, args_run_demo = dd_argparse.process_pipeline_args(args)

    # Initialize demo
    _ = kwargs_init_pipeline.pop('guidance_scale')
    _ = kwargs_init_pipeline.pop('denoising_steps')
    demo = StableCascadeDemoPipeline(
        args.prior_denoising_steps,
        args.decoder_denoising_steps,
        args.prior_guidance_scale,
        args.decoder_guidance_scale,
        args.lite,
        **kwargs_init_pipeline
    )

    # Load TensorRT engines and pytorch modules
    demo.loadEngines(
        args.framework_model_dir,
        args.onnx_dir,
        args.engine_dir,
        **kwargs_load_engine,
    )

    # Load resources
    _, shared_device_memory = cudart.cudaMalloc(demo.get_max_device_memory())
    demo.activateEngines(shared_device_memory)
    demo.loadResources(args.height, args.width, args.batch_size, args.seed)

    # Run inference
    demo.run(*args_run_demo)

    demo.teardown()