# # 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. # import gc import inspect import json import os import pathlib import sys import time from hashlib import md5 from typing import List, Optional import modelopt.torch.opt as mto import modelopt.torch.quantization as mtq import numpy as np import nvtx import tensorrt as trt import torch from cuda.bindings import runtime as cudart from diffusers import ( DDIMScheduler, DDPMScheduler, DDPMWuerstchenScheduler, EulerAncestralDiscreteScheduler, EulerDiscreteScheduler, LCMScheduler, LMSDiscreteScheduler, PNDMScheduler, UniPCMultistepScheduler, ) import demo_diffusion.engine as engine_module import demo_diffusion.image as image_module from demo_diffusion.model import ( CLIPModel, CLIPWithProjModel, SDLoraLoader, UNet2DConditionControlNetModel, UNetModel, UNetXLModel, UNetXLModelControlNet, VAEEncoderModel, VAEModel, get_clip_embedding_dim, make_scheduler, make_tokenizer, merge_loras, unload_torch_model, ) from demo_diffusion.pipeline.calibrate import load_calib_prompts from demo_diffusion.pipeline.type import PIPELINE_TYPE from demo_diffusion.utils_modelopt import ( SD_FP8_FP16_DEFAULT_CONFIG, SD_FP8_FP32_DEFAULT_CONFIG, check_lora, filter_func, generate_fp8_scales, get_int8_config, quantize_lvl, set_fmha, ) TRT_LOGGER = trt.Logger(trt.Logger.ERROR) class StableDiffusionPipeline: SCHEDULER_DEFAULTS = { "1.4": "PNDM", "dreamshaper-7": "PNDM", "xl-1.0" : "Euler", "xl-turbo": "EulerA", "svd-xt-1.1": "Euler", "cascade": "DDPMWuerstchen" } """ Application showcasing the acceleration of Stable Diffusion pipelines using NVidia TensorRT. """ def __init__( self, version='1.4', pipeline_type=PIPELINE_TYPE.TXT2IMG, max_batch_size=16, denoising_steps=30, scheduler=None, guidance_scale=7.5, device='cuda', output_dir='.', hf_token=None, verbose=False, nvtx_profile=False, use_cuda_graph=False, vae_scaling_factor=0.18215, framework_model_dir='pytorch_model', controlnets=None, lora_scale: float = 1.0, lora_weight: Optional[List[float]] = None, lora_path: Optional[List[str]] = None, return_latents=False, torch_inference='', ): """ Initializes the Diffusion pipeline. Args: version (str): The version of the pipeline. Should be one of [1.4, SDXL] pipeline_type (PIPELINE_TYPE): Type of current pipeline. max_batch_size (int): Maximum batch size for dynamic batch engine. denoising_steps (int): The number of denoising steps. More denoising steps usually lead to a higher quality image at the expense of slower inference. scheduler (str): The scheduler to guide the denoising process. Must be one of [DDIM, DPM, EulerA, Euler, LCM, LMSD, PNDM]. 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. nvtx_profile (bool): Insert NVTX profiling markers. use_cuda_graph (bool): Use CUDA graph to capture engine execution and then launch inference vae_scaling_factor (float): VAE scaling factor framework_model_dir (str): cache directory for framework checkpoints controlnets (str): Which ControlNet/ControlNets to use. return_latents (bool): Skip decoding the image and return latents instead. torch_inference (str): Run inference with PyTorch (using specified compilation mode) instead of TensorRT. """ self.denoising_steps = denoising_steps self.guidance_scale = guidance_scale self.do_classifier_free_guidance = (guidance_scale > 1.0) self.vae_scaling_factor = vae_scaling_factor self.max_batch_size = max_batch_size self.framework_model_dir = framework_model_dir self.output_dir = output_dir for directory in [self.framework_model_dir, self.output_dir]: if not os.path.exists(directory): print(f"[I] Create directory: {directory}") pathlib.Path(directory).mkdir(parents=True) self.hf_token = hf_token self.device = device self.verbose = verbose self.nvtx_profile = nvtx_profile self.version = version self.controlnets = controlnets # Pipeline type self.pipeline_type = pipeline_type if self.pipeline_type.is_txt2img(): self.stages = ['clip','unet','vae'] elif self.pipeline_type.is_img2img(): self.stages = ['vae_encoder', 'clip','unet','vae'] elif self.pipeline_type.is_sd_xl_base(): self.stages = ['clip', 'clip2', 'unetxl'] if not return_latents: self.stages.append('vae') elif self.pipeline_type.is_sd_xl_refiner(): self.stages = ['clip2', 'unetxl', 'vae'] elif self.pipeline_type.is_img2vid(): self.stages = ['clip-vis', 'clip-imgfe', 'unet-temp', 'vae-temp'] elif self.pipeline_type.is_cascade_prior(): self.stages = ['clip', 'unet'] elif self.pipeline_type.is_cascade_decoder(): self.stages = ['clip', 'unet', 'vqgan'] else: raise ValueError(f"Unsupported pipeline {self.pipeline_type.name}.") self.return_latents = return_latents if not scheduler: scheduler = 'UniPC' if self.pipeline_type.is_controlnet() else self.SCHEDULER_DEFAULTS.get(version, 'DDIM') print(f"[I] Autoselected scheduler: {scheduler}") scheduler_class_map = { "DDIM" : DDIMScheduler, "DDPM" : DDPMScheduler, "EulerA" : EulerAncestralDiscreteScheduler, "Euler" : EulerDiscreteScheduler, "LCM" : LCMScheduler, "LMSD" : LMSDiscreteScheduler, "PNDM" : PNDMScheduler, "UniPC" : UniPCMultistepScheduler, "DDPMWuerstchen" : DDPMWuerstchenScheduler, } try: scheduler_class = scheduler_class_map[scheduler] except KeyError: raise ValueError(f"Unsupported scheduler {scheduler}. Should be one of {list(scheduler_class.keys())}.") self.scheduler = make_scheduler(scheduler_class, version, pipeline_type, hf_token, framework_model_dir) self.config = {} if self.pipeline_type.is_sd_xl(): self.config['clip_hidden_states'] = True self.torch_inference = torch_inference if self.torch_inference: torch._inductor.config.conv_1x1_as_mm = True torch._inductor.config.coordinate_descent_tuning = True torch._inductor.config.epilogue_fusion = False torch._inductor.config.coordinate_descent_check_all_directions = True self.use_cuda_graph = use_cuda_graph # initialized in loadEngines() self.models = {} self.torch_models = {} self.engine = {} self.shared_device_memory = None # initialize lora loader and scales self.lora_loader = None self.lora_weights = dict() if lora_path: self.lora_loader = SDLoraLoader(lora_path, lora_weight, lora_scale) assert len(lora_path) == len(lora_weight) for i, path in enumerate(lora_path): self.lora_weights[path] = lora_weight[i] # initialized in loadResources() self.events = {} self.generator = None self.markers = {} self.seed = None self.stream = None self.tokenizer = None def loadResources(self, image_height, image_width, batch_size, seed): # Initialize noise generator if seed: self.seed = seed self.generator = torch.Generator(device="cuda").manual_seed(seed) # Create CUDA events and stream for stage in self.stages: self.events[stage] = [cudart.cudaEventCreate()[1], cudart.cudaEventCreate()[1]] self.stream = cudart.cudaStreamCreate()[1] # Allocate TensorRT I/O buffers if not self.torch_inference: for model_name, obj in self.models.items(): self.engine[model_name].allocate_buffers(shape_dict=obj.get_shape_dict(batch_size, image_height, image_width), device=self.device) def teardown(self): for e in self.events.values(): cudart.cudaEventDestroy(e[0]) cudart.cudaEventDestroy(e[1]) for engine in self.engine.values(): del engine if self.shared_device_memory: cudart.cudaFree(self.shared_device_memory) for torch_model in self.torch_models.values(): del torch_model cudart.cudaStreamDestroy(self.stream) del self.stream def cachedModelName(self, model_name): return model_name def getOnnxPath(self, model_name, onnx_dir, opt=True, suffix=''): onnx_model_dir = os.path.join(onnx_dir, self.cachedModelName(model_name)+suffix+('.opt' if opt else '')) os.makedirs(onnx_model_dir, exist_ok=True) return os.path.join(onnx_model_dir, 'model.onnx') def getEnginePath(self, model_name, engine_dir, enable_refit=False, suffix=''): return os.path.join(engine_dir, self.cachedModelName(model_name)+suffix+('.refit' if enable_refit else '')+'.trt'+trt.__version__+'.plan') def getWeightsMapPath(self, model_name, onnx_dir): onnx_model_dir = os.path.join(onnx_dir, self.cachedModelName(model_name)+'.opt') os.makedirs(onnx_model_dir, exist_ok=True) return os.path.join(onnx_model_dir, 'weights_map.json') def getRefitNodesPath(self, model_name, onnx_dir, suffix=''): onnx_model_dir = os.path.join(onnx_dir, self.cachedModelName(model_name)+'.opt') os.makedirs(onnx_model_dir, exist_ok=True) return os.path.join(onnx_model_dir, 'refit'+suffix+'.json') def getStateDictPath(self, model_name, onnx_dir, suffix=''): onnx_model_dir = os.path.join(onnx_dir, self.cachedModelName(model_name)+suffix) os.makedirs(onnx_model_dir, exist_ok=True) return os.path.join(onnx_model_dir, 'state_dict.pt') def initializeModels(self, framework_model_dir, int8, fp8): # Load text tokenizer(s) if not self.pipeline_type.is_sd_xl_refiner(): self.tokenizer = make_tokenizer(self.version, self.pipeline_type, self.hf_token, framework_model_dir) if self.pipeline_type.is_sd_xl(): self.tokenizer2 = make_tokenizer(self.version, self.pipeline_type, self.hf_token, framework_model_dir, subfolder='tokenizer_2') # Load pipeline models models_args = {'version': self.version, 'pipeline': self.pipeline_type, 'device': self.device, 'hf_token': self.hf_token, 'verbose': self.verbose, 'framework_model_dir': framework_model_dir, 'max_batch_size': self.max_batch_size} if 'clip' in self.stages: subfolder = 'text_encoder' self.models['clip'] = CLIPModel(**models_args, fp16=True, embedding_dim=get_clip_embedding_dim(self.version, self.pipeline_type), output_hidden_states=self.config.get('clip_hidden_states', False), subfolder=subfolder) if 'clip2' in self.stages: subfolder = 'text_encoder_2' self.models['clip2'] = CLIPWithProjModel(**models_args, fp16=True, output_hidden_states=self.config.get('clip_hidden_states', False), subfolder=subfolder) if 'unet' in self.stages: self.models['unet'] = UNetModel(**models_args, fp16=True, int8=int8, fp8=fp8, controlnets=self.controlnets, do_classifier_free_guidance=self.do_classifier_free_guidance) if 'unetxl' in self.stages: if not self.controlnets: self.models["unetxl"] = UNetXLModel( **models_args, fp16=True, int8=int8, fp8=fp8, do_classifier_free_guidance=self.do_classifier_free_guidance, ) else: self.models["unetxl"] = UNetXLModelControlNet( **models_args, fp16=True, int8=int8, fp8=fp8, controlnets=self.controlnets, do_classifier_free_guidance=self.do_classifier_free_guidance, ) vae_fp16 = not self.pipeline_type.is_sd_xl() if 'vae' in self.stages: self.models['vae'] = VAEModel(**models_args, fp16=vae_fp16, tf32=True) if 'vae_encoder' in self.stages: self.models['vae_encoder'] = VAEEncoderModel(**models_args, fp16=vae_fp16) def loadEngines( self, engine_dir, framework_model_dir, onnx_dir, onnx_opset, opt_batch_size, opt_image_height, opt_image_width, optimization_level=3, static_batch=False, static_shape=True, enable_refit=False, enable_all_tactics=False, timing_cache=None, int8=False, fp8=False, quantization_level=2.5, quantization_percentile=1.0, quantization_alpha=0.8, calibration_size=32, calib_batch_size=2, **_kwargs, ): """ Build and load engines for TensorRT accelerated inference. Export ONNX models first, if applicable. Args: engine_dir (str): Directory to store the TensorRT engines. framework_model_dir (str): Directory to store the framework model ckpt. onnx_dir (str): Directory to store 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. opt_image_height (int): Image height to optimize for during engine building. Must be a multiple of 8. opt_image_width (int): Image width to optimize for during engine building. Must be a multiple of 8. optimization_level (int): Optimization level to build the TensorRT engine with. static_batch (bool): Build engine only for specified opt_batch_size. static_shape (bool): Build engine only for specified opt_image_height & opt_image_width. Default = True. enable_refit (bool): Build engines with refit option enabled. enable_all_tactics (bool): Enable all tactic sources during TensorRT engine builds. timing_cache (str): Path to the timing cache to speed up TensorRT build. int8 (bool): Whether to quantize to int8 format or not (SDXL, SD15 and SD21 only). fp8 (bool): Whether to quantize to fp8 format or not (SDXL, SD15 and SD21 only). quantization_level (float): Controls which layers to quantize. 1: CNN, 2: CNN+FFN, 2.5: CNN+FFN+QKV, 3: CNN+FC quantization_percentile (float): Control quantization scaling factors (amax) collecting range, where the minimum amax in range(n_steps * percentile) will be collected. Recommendation: 1.0 quantization_alpha (float): The alpha parameter for SmoothQuant quantization used for linear layers. Recommendation: 0.8 for SDXL calibration_size (int): The number of steps to use for calibrating the model for quantization. Recommendation: 32, 64, 128 for SDXL calib_batch_size (int): The batch size to use for calibration. Defaults to 2. """ # Create directories if missing for directory in [engine_dir, onnx_dir]: if not os.path.exists(directory): print(f"[I] Create directory: {directory}") pathlib.Path(directory).mkdir(parents=True) # Initialize models self.initializeModels(framework_model_dir, int8, fp8) # Configure pipeline models to load model_names = self.models.keys() lora_suffix = '-'+'-'.join([str(md5(path.encode('utf-8')).hexdigest())+'-'+('%.2f' % self.lora_weights[path])+'-'+('%.2f' % self.lora_loader.scale) for path in sorted(self.lora_loader.paths)]) if self.lora_loader else '' # Enable refit and LoRA merging only for UNet & UNetXL for now do_engine_refit = dict(zip(model_names, [not self.pipeline_type.is_sd_xl_refiner() and enable_refit and model_name.startswith('unet') for model_name in model_names])) do_lora_merge = dict(zip(model_names, [not enable_refit and self.lora_loader and model_name.startswith('unet') for model_name in model_names])) # Torch fallback for VAE if specified torch_fallback = dict(zip(model_names, [self.torch_inference for model_name in model_names])) model_suffix = dict(zip(model_names, [lora_suffix if do_lora_merge[model_name] else '' for model_name in model_names])) use_int8 = dict.fromkeys(model_names, False) use_fp8 = dict.fromkeys(model_names, False) if int8: assert self.pipeline_type.is_sd_xl_base() or self.version == "1.4", "int8 quantization only supported for SDXL and SD1.4 pipeline" model_name = 'unetxl' if self.pipeline_type.is_sd_xl() else 'unet' use_int8[model_name] = True model_suffix[model_name] += f"-int8.l{quantization_level}.bs2.s{self.denoising_steps}.c{calibration_size}.p{quantization_percentile}.a{quantization_alpha}" elif fp8: assert self.pipeline_type.is_sd_xl() or self.version == "1.4", "fp8 quantization only supported for SDXL and SD1.4 pipeline" model_name = 'unetxl' if self.pipeline_type.is_sd_xl() else 'unet' use_fp8[model_name] = True model_suffix[model_name] += f"-fp8.l{quantization_level}.bs2.s{self.denoising_steps}.c{calibration_size}.p{quantization_percentile}.a{quantization_alpha}" onnx_path = dict(zip(model_names, [self.getOnnxPath(model_name, onnx_dir, opt=False, suffix=model_suffix[model_name]) for model_name in model_names])) onnx_opt_path = dict(zip(model_names, [self.getOnnxPath(model_name, onnx_dir, suffix=model_suffix[model_name]) for model_name in model_names])) engine_path = dict(zip(model_names, [self.getEnginePath(model_name, engine_dir, do_engine_refit[model_name], suffix=model_suffix[model_name]) for model_name in model_names])) weights_map_path = dict(zip(model_names, [(self.getWeightsMapPath(model_name, onnx_dir) if do_engine_refit[model_name] else None) for model_name in model_names])) for model_name, obj in self.models.items(): if torch_fallback[model_name]: continue # Export models to ONNX and save weights name mapping do_export_onnx = not os.path.exists(engine_path[model_name]) and not os.path.exists(onnx_opt_path[model_name]) do_export_weights_map = weights_map_path[model_name] and not os.path.exists(weights_map_path[model_name]) if do_export_onnx or do_export_weights_map: # Non-quantized ONNX export if not use_int8[model_name] and not use_fp8[model_name]: obj.export_onnx(onnx_path[model_name], onnx_opt_path[model_name], onnx_opset, opt_image_height, opt_image_width, enable_lora_merge=do_lora_merge[model_name], static_shape=static_shape, lora_loader=self.lora_loader) else: pipeline = obj.get_pipeline() if self.pipeline_type.is_controlnet(): model = UNet2DConditionControlNetModel(pipeline.unet, pipeline.controlnet.nets) else: model = pipeline.unet if use_fp8[model_name] and quantization_level == 4.0: set_fmha(model) state_dict_path = self.getStateDictPath(model_name, onnx_dir, suffix=model_suffix[model_name]) if not os.path.exists(state_dict_path): print(f"[I] Calibrated weights not found, generating {state_dict_path}") root_dir = os.path.dirname(os.path.abspath(sys.modules["__main__"].__file__)) calibration_file = os.path.join(root_dir, "calibration_data", "calibration-prompts.txt") calibration_prompts = load_calib_prompts(calib_batch_size, calibration_file) if self.pipeline_type.is_controlnet(): calibration_image_canny = image_module.download_image( "https://huggingface.co/diffusers/controlnet-canny-sdxl-1.0/resolve/main/out_bird.png" ) # "out_bird.png" has 5 images combined in a row. We pick the first image which is the input image. calibration_image_canny = calibration_image_canny.crop( (0, 0, calibration_image_canny.width / 5, calibration_image_canny.height) ) calibration_images = [calibration_image_canny] # TODO check size > calibration_size def do_calibrate(pipeline, calibration_prompts, **kwargs): for i_th, prompts in enumerate(calibration_prompts): if i_th >= kwargs["calib_size"]: return pipeline_call_kwargs = { "prompt": prompts, "num_inference_steps": kwargs["n_steps"], "negative_prompt": [ "normal quality, low quality, worst quality, low res, blurry, nsfw, nude" ] * len(prompts), } if self.pipeline_type.is_controlnet(): pipeline_call_kwargs["image"] = calibration_images pipeline(**pipeline_call_kwargs).images def forward_loop(model): if self.pipeline_type.is_controlnet(): pipeline.unet = model.unet pipeline.controlnet.nets = model.controlnets else: pipeline.unet = model do_calibrate( pipeline=pipeline, calibration_prompts=calibration_prompts, calib_size=calibration_size // calib_batch_size, n_steps=self.denoising_steps, ) print(f"[I] Performing calibration for {calibration_size} steps.") if use_int8[model_name]: quant_config = get_int8_config( model, quantization_level, quantization_alpha, quantization_percentile, self.denoising_steps ) elif use_fp8[model_name]: quant_config = SD_FP8_FP16_DEFAULT_CONFIG # Handle LoRA if do_lora_merge[model_name]: assert self.lora_loader is not None model = merge_loras(model, self.lora_loader) check_lora(model) mtq.quantize(model, quant_config, forward_loop) mto.save(model, state_dict_path) else: mto.restore(model, state_dict_path) print(f"[I] Generating quantized ONNX model: {onnx_opt_path[model_name]}") if not os.path.exists(onnx_path[model_name]): quantize_lvl(self.version, model, quantization_level) mtq.disable_quantizer(model, filter_func) if use_fp8[model_name]: generate_fp8_scales(model) else: model = None obj.export_onnx(onnx_path[model_name], onnx_opt_path[model_name], onnx_opset, opt_image_height, opt_image_width, custom_model=model, static_shape=static_shape) # FIXME do_export_weights_map needs ONNX graph if do_export_weights_map: print(f"[I] Saving weights map: {weights_map_path[model_name]}") obj.export_weights_map(onnx_opt_path[model_name], weights_map_path[model_name]) # Release temp GPU memory during onnx export to avoid OOM. gc.collect() torch.cuda.empty_cache() # Build TensorRT engines for model_name, obj in self.models.items(): if torch_fallback[model_name]: continue engine = engine_module.Engine(engine_path[model_name]) if not os.path.exists(engine_path[model_name]): update_output_names = obj.get_output_names() + obj.extra_output_names if obj.extra_output_names else None # TF32 can be enabled for all precisions (including INT8/FP8) tf32amp = obj.tf32 precision_constraints = 'none' engine.build(onnx_opt_path[model_name], tf32=tf32amp, input_profile=obj.get_input_profile( opt_batch_size, opt_image_height, opt_image_width, static_batch=static_batch, static_shape=static_shape ), enable_refit=do_engine_refit[model_name], enable_all_tactics=enable_all_tactics, timing_cache=timing_cache, update_output_names=update_output_names, verbose=self.verbose, builder_optimization_level=optimization_level, precision_constraints=precision_constraints, ) self.engine[model_name] = engine # Load TensorRT engines for model_name, obj in self.models.items(): if torch_fallback[model_name]: continue self.engine[model_name].load() if do_engine_refit[model_name] and self.lora_loader: assert weights_map_path[model_name] with open(weights_map_path[model_name], 'r') as fp_wts: print(f"[I] Loading weights map: {weights_map_path[model_name]} ") [weights_name_mapping, weights_shape_mapping] = json.load(fp_wts) refit_weights_path = self.getRefitNodesPath(model_name, engine_dir, suffix=lora_suffix) if not os.path.exists(refit_weights_path): print(f"[I] Saving refit weights: {refit_weights_path}") model = merge_loras(obj.get_model(), self.lora_loader) refit_weights, updated_weight_names = engine_module.get_refit_weights( model.state_dict(), onnx_opt_path[model_name], weights_name_mapping, weights_shape_mapping ) torch.save((refit_weights, updated_weight_names), refit_weights_path) unload_torch_model(model) else: print(f"[I] Loading refit weights: {refit_weights_path}") refit_weights, updated_weight_names = torch.load(refit_weights_path) self.engine[model_name].refit(refit_weights, updated_weight_names) # Load torch models for model_name, obj in self.models.items(): if torch_fallback[model_name]: self.torch_models[model_name] = obj.get_model(torch_inference=self.torch_inference) # Release temp GPU memory during onnx export to avoid OOM. gc.collect() torch.cuda.empty_cache() def calculateMaxDeviceMemory(self): max_device_memory = 0 for model_name, engine in self.engine.items(): max_device_memory = max(max_device_memory, engine.engine.device_memory_size_v2) return max_device_memory def activateEngines(self, shared_device_memory=None): if shared_device_memory is None: max_device_memory = self.calculateMaxDeviceMemory() _, shared_device_memory = cudart.cudaMalloc(max_device_memory) self.shared_device_memory = shared_device_memory # Load and activate TensorRT engines for engine in self.engine.values(): engine.activate(device_memory=self.shared_device_memory) def runEngine(self, model_name, feed_dict): engine = self.engine[model_name] return engine.infer(feed_dict, self.stream, use_cuda_graph=self.use_cuda_graph) def initialize_latents(self, batch_size, unet_channels, latent_height, latent_width, latents_dtype=torch.float32): latents_dtype = latents_dtype # text_embeddings.dtype latents_shape = (batch_size, unet_channels, latent_height, latent_width) latents = torch.randn(latents_shape, device=self.device, dtype=latents_dtype, generator=self.generator) # Scale the initial noise by the standard deviation required by the scheduler latents = latents * self.scheduler.init_noise_sigma return latents def profile_start(self, name, color='blue'): if self.nvtx_profile: self.markers[name] = nvtx.start_range(message=name, color=color) if name in self.events: cudart.cudaEventRecord(self.events[name][0], 0) def profile_stop(self, name): if name in self.events: cudart.cudaEventRecord(self.events[name][1], 0) if self.nvtx_profile: nvtx.end_range(self.markers[name]) def preprocess_images(self, batch_size, images=()): if not images: return () self.profile_start('preprocess', color='pink') input_images=[] for image in images: image = image.to(self.device).float() if image.shape[0] != batch_size: image = image.repeat(batch_size, 1, 1, 1) input_images.append(image) self.profile_stop('preprocess') return tuple(input_images) def preprocess_controlnet_images(self, batch_size, images=None): ''' images: List of PIL.Image.Image ''' if images is None: return None self.profile_start('preprocess', color='pink') images = [(np.array(i.convert("RGB")).astype(np.float32) / 255.0)[..., None].transpose(3, 2, 0, 1).repeat(batch_size, axis=0) for i in images] # do_classifier_free_guidance images = [torch.cat([torch.from_numpy(i).to(self.device).float()] * 2) for i in images] images = torch.cat([image[None, ...] for image in images], dim=0) self.profile_stop('preprocess') return images def encode_prompt(self, prompt, negative_prompt, encoder='clip', pooled_outputs=False, output_hidden_states=False): self.profile_start(encoder, color='green') tokenizer = self.tokenizer2 if encoder == 'clip2' else self.tokenizer def tokenize(prompt, output_hidden_states): text_input_ids = tokenizer( prompt, padding="max_length", max_length=tokenizer.model_max_length, truncation=True, return_tensors="pt", ).input_ids.type(torch.int32).to(self.device) text_hidden_states = None if self.torch_inference: outputs = self.torch_models[encoder](text_input_ids, output_hidden_states=output_hidden_states) text_embeddings = outputs[0].clone() if output_hidden_states: text_hidden_states = outputs['hidden_states'][-2].clone() else: # NOTE: output tensor for CLIP must be cloned because it will be overwritten when called again for negative prompt outputs = self.runEngine(encoder, {'input_ids': text_input_ids}) text_embeddings = outputs['text_embeddings'].clone() if output_hidden_states: text_hidden_states = outputs['hidden_states'].clone() return text_embeddings, text_hidden_states # Tokenize prompt text_embeddings, text_hidden_states = tokenize(prompt, output_hidden_states) if self.do_classifier_free_guidance: # Tokenize negative prompt uncond_embeddings, uncond_hidden_states = tokenize(negative_prompt, output_hidden_states) # 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]).to(dtype=torch.float16) if pooled_outputs: pooled_output = text_embeddings if output_hidden_states: text_embeddings = torch.cat([uncond_hidden_states, text_hidden_states]).to(dtype=torch.float16) if self.do_classifier_free_guidance else text_hidden_states self.profile_stop(encoder) if pooled_outputs: return text_embeddings, pooled_output return text_embeddings # from diffusers (get_timesteps) def get_timesteps(self, num_inference_steps, strength, denoising_start=None): # get the original timestep using init_timestep if denoising_start is None: init_timestep = min(int(num_inference_steps * strength), num_inference_steps) t_start = max(num_inference_steps - init_timestep, 0) else: t_start = 0 timesteps = self.scheduler.timesteps[t_start * self.scheduler.order :] # Strength is irrelevant if we directly request a timestep to start at; # that is, strength is determined by the denoising_start instead. if denoising_start is not None: discrete_timestep_cutoff = int( round( self.scheduler.config.num_train_timesteps - (denoising_start * self.scheduler.config.num_train_timesteps) ) ) num_inference_steps = (timesteps < discrete_timestep_cutoff).sum().item() if self.scheduler.order == 2 and num_inference_steps % 2 == 0: # if the scheduler is a 2nd order scheduler we might have to do +1 # because `num_inference_steps` might be even given that every timestep # (except the highest one) is duplicated. If `num_inference_steps` is even it would # mean that we cut the timesteps in the middle of the denoising step # (between 1st and 2nd devirative) which leads to incorrect results. By adding 1 # we ensure that the denoising process always ends after the 2nd derivate step of the scheduler num_inference_steps = num_inference_steps + 1 # because t_n+1 >= t_n, we slice the timesteps starting from the end timesteps = timesteps[-num_inference_steps:] return timesteps, num_inference_steps return timesteps, num_inference_steps - t_start def denoise_latent(self, latents, text_embeddings, denoiser='unet', timesteps=None, step_offset=0, mask=None, masked_image_latents=None, image_guidance=1.5, controlnet_imgs=None, controlnet_scales=None, text_embeds=None, time_ids=None): assert image_guidance > 1.0, "Image guidance has to be > 1.0" controlnet_imgs = self.preprocess_controlnet_images(latents.shape[0], controlnet_imgs) do_autocast = self.torch_inference != '' and self.models[denoiser].fp16 with torch.autocast('cuda', enabled=do_autocast): self.profile_start(denoiser, color='blue') for step_index, timestep in enumerate(timesteps): # Expand the latents if we are doing classifier free guidance latent_model_input = torch.cat([latents] * 2) if self.do_classifier_free_guidance else latents latent_model_input = self.scheduler.scale_model_input(latent_model_input, timestep) if isinstance(mask, torch.Tensor): latent_model_input = torch.cat([latent_model_input, mask, masked_image_latents], dim=1) # Predict the noise residual if self.torch_inference: params = {"sample": latent_model_input, "timestep": timestep, "encoder_hidden_states": text_embeddings} if controlnet_imgs is not None: params.update({"images": controlnet_imgs, "controlnet_scales": controlnet_scales}) added_cond_kwargs = {} if text_embeds != None: added_cond_kwargs.update({'text_embeds': text_embeds}) if time_ids != None: added_cond_kwargs.update({'time_ids': time_ids}) if text_embeds != None or time_ids != None: params.update({'added_cond_kwargs': added_cond_kwargs}) noise_pred = self.torch_models[denoiser](**params)["sample"] else: timestep_float = timestep.float() if timestep.dtype != torch.float32 else timestep params = {"sample": latent_model_input, "timestep": timestep_float, "encoder_hidden_states": text_embeddings} if controlnet_imgs is not None: params.update({"images": controlnet_imgs, "controlnet_scales": controlnet_scales}) if text_embeds != None: params.update({'text_embeds': text_embeds}) if time_ids != None: params.update({'time_ids': time_ids}) noise_pred = self.runEngine(denoiser, params)['latent'] # Perform guidance if self.do_classifier_free_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) # from diffusers (prepare_extra_step_kwargs) extra_step_kwargs = {} if "eta" in set(inspect.signature(self.scheduler.step).parameters.keys()): # TODO: configurable eta eta = 0.0 extra_step_kwargs["eta"] = eta if "generator" in set(inspect.signature(self.scheduler.step).parameters.keys()): extra_step_kwargs["generator"] = self.generator latents = self.scheduler.step(noise_pred, timestep, latents, **extra_step_kwargs, return_dict=False)[0] latents = 1. / self.vae_scaling_factor * latents latents = latents.to(dtype=torch.float32) self.profile_stop(denoiser) return latents def encode_image(self, input_image): self.profile_start('vae_encoder', color='red') cast_to = torch.float16 if self.models['vae_encoder'].fp16 else torch.bfloat16 if self.models['vae_encoder'].bf16 else torch.float32 input_image = input_image.to(dtype=cast_to) if self.torch_inference: image_latents = self.torch_models['vae_encoder'](input_image) else: image_latents = self.runEngine('vae_encoder', {'images': input_image})['latent'] image_latents = self.vae_scaling_factor * image_latents self.profile_stop('vae_encoder') return image_latents def decode_latent(self, latents): self.profile_start('vae', color='red') cast_to = torch.float16 if self.models['vae'].fp16 else torch.bfloat16 if self.models['vae'].bf16 else torch.float32 latents = latents.to(dtype=cast_to) if self.torch_inference: images = self.torch_models['vae'](latents, return_dict=False)[0] else: images = self.runEngine('vae', {'latent': latents})['images'] self.profile_stop('vae') return images def print_summary(self, denoising_steps, walltime_ms, batch_size): print('|-----------------|--------------|') print('| {:^15} | {:^12} |'.format('Module', 'Latency')) print('|-----------------|--------------|') for stage in self.stages: stage_name = stage if "unet" in stage: if self.pipeline_type.is_controlnet(): stage_name += '+cnet' stage_name += ' x ' + str(denoising_steps) print( "| {:^15} | {:>9.2f} ms |".format( stage_name, cudart.cudaEventElapsedTime(self.events[stage][0], self.events[stage][1])[1], ) ) print('|-----------------|--------------|') print('| {:^15} | {:>9.2f} ms |'.format('Pipeline', walltime_ms)) print('|-----------------|--------------|') print('Throughput: {:.5f} image/s'.format(batch_size*1000./walltime_ms)) def save_image(self, images, pipeline, prompt, seed): # Save image image_name_prefix = pipeline+''.join(set(['-'+prompt[i].replace(' ','_')[:10] for i in range(len(prompt))]))+'-'+str(seed)+'-' image_name_suffix = 'torch' if self.torch_inference else 'trt' image_module.save_image(images, self.output_dir, image_name_prefix, image_name_suffix) def infer( self, prompt, negative_prompt, image_height, image_width, input_image=None, image_strength=0.75, controlnet_scales=None, aesthetic_score=6.0, negative_aesthetic_score=2.5, warmup=False, verbose=False, save_image=True, ): """ 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. image_height (int): Height (in pixels) of the image to be generated. Must be a multiple of 8. image_width (int): Width (in pixels) of the image to be generated. Must be a multiple of 8. input_image (image): Input image used to initialize the latents. image_strength (float): Strength of transformation applied to input_image. Must be between 0 and 1. controlnet_scales (torch.Tensor) A tensor which containes ControlNet scales, essential for multi ControlNet. Must be equal to number of Controlnets. warmup (bool): Indicate if this is a warmup run. verbose (bool): Verbose in logging save_image (bool): Save the generated image (if applicable) """ assert len(prompt) == len(negative_prompt) batch_size = len(prompt) # Spatial dimensions of latent tensor latent_height = image_height // 8 latent_width = image_width // 8 if self.generator and self.seed: self.generator.manual_seed(self.seed) num_inference_steps = self.denoising_steps with torch.inference_mode(), trt.Runtime(TRT_LOGGER): torch.cuda.synchronize() e2e_tic = time.perf_counter() # TODO: support custom timesteps timesteps = None if timesteps is not None: if "timesteps" not in set(inspect.signature(self.scheduler.set_timesteps).parameters.keys()): raise ValueError( f"The current scheduler class {self.scheduler.__class__}'s `set_timesteps` does not support custom" f" timestep schedules. Please check whether you are using the correct scheduler." ) self.scheduler.set_timesteps(timesteps=timesteps, device=self.device) assert self.denoising_steps == len(self.scheduler.timesteps) else: self.scheduler.set_timesteps(self.denoising_steps, device=self.device) timesteps = self.scheduler.timesteps.to(self.device) denoise_kwargs = {} if not (self.pipeline_type.is_img2img() or self.pipeline_type.is_sd_xl_refiner()): # Initialize latents latents = self.initialize_latents(batch_size=batch_size, unet_channels=4, latent_height=latent_height, latent_width=latent_width) if self.pipeline_type.is_controlnet(): denoise_kwargs.update({'controlnet_imgs': input_image, 'controlnet_scales': controlnet_scales}) # Pre-process and VAE encode input image if self.pipeline_type.is_img2img() or self.pipeline_type.is_sd_xl_refiner(): assert input_image != None # Initialize timesteps and pre-process input image timesteps, num_inference_steps = self.get_timesteps(self.denoising_steps, image_strength) denoise_kwargs.update({'timesteps': timesteps}) if self.pipeline_type.is_img2img() or self.pipeline_type.is_sd_xl_refiner(): latent_timestep = timesteps[:1].repeat(batch_size) input_image = self.preprocess_images(batch_size, (input_image,))[0] # Encode if not a latent image_latents = input_image if input_image.shape[1] == 4 else self.encode_image(input_image) # Add noise to latents using timesteps noise = torch.randn(image_latents.shape, generator=self.generator, device=self.device, dtype=torch.float32) latents = self.scheduler.add_noise(image_latents, noise, latent_timestep) # CLIP text encoder(s) if self.pipeline_type.is_sd_xl(): text_embeddings2, pooled_embeddings2 = self.encode_prompt(prompt, negative_prompt, encoder='clip2', pooled_outputs=True, output_hidden_states=True) # Merge text embeddings if self.pipeline_type.is_sd_xl_base(): text_embeddings = self.encode_prompt(prompt, negative_prompt, output_hidden_states=True) text_embeddings = torch.cat([text_embeddings, text_embeddings2], dim=-1) else: text_embeddings = text_embeddings2 # Time embeddings def _get_add_time_ids(original_size, crops_coords_top_left, target_size, dtype, aesthetic_score=None, negative_aesthetic_score=None): if self.pipeline_type.is_sd_xl_refiner(): #self.requires_aesthetics_score: add_time_ids = list(original_size + crops_coords_top_left + (aesthetic_score,)) if self.do_classifier_free_guidance: add_neg_time_ids = list(original_size + crops_coords_top_left + (negative_aesthetic_score,)) else: add_time_ids = list(original_size + crops_coords_top_left + target_size) if self.do_classifier_free_guidance: add_neg_time_ids = list(original_size + crops_coords_top_left + target_size) add_time_ids = torch.tensor([add_time_ids], dtype=dtype, device=self.device) if self.do_classifier_free_guidance: add_neg_time_ids = torch.tensor([add_neg_time_ids], dtype=dtype, device=self.device) add_time_ids = torch.cat([add_neg_time_ids, add_time_ids], dim=0) return add_time_ids original_size = (image_height, image_width) crops_coords_top_left = (0, 0) target_size = (image_height, image_width) if self.pipeline_type.is_sd_xl_refiner(): add_time_ids = _get_add_time_ids( original_size, crops_coords_top_left, target_size, dtype=text_embeddings.dtype, aesthetic_score=aesthetic_score, negative_aesthetic_score=negative_aesthetic_score ) else: add_time_ids = _get_add_time_ids( original_size, crops_coords_top_left, target_size, dtype=text_embeddings.dtype ) add_time_ids = add_time_ids.repeat(batch_size, 1) denoise_kwargs.update({'text_embeds': pooled_embeddings2, 'time_ids': add_time_ids}) else: text_embeddings = self.encode_prompt(prompt, negative_prompt) # UNet denoiser + (optional) ControlNet(s) denoiser = 'unetxl' if self.pipeline_type.is_sd_xl() else 'unet' latents = self.denoise_latent(latents, text_embeddings, denoiser=denoiser, **denoise_kwargs) # VAE decode latent (if applicable) if self.return_latents: latents = latents * self.vae_scaling_factor else: images = self.decode_latent(latents) torch.cuda.synchronize() e2e_toc = time.perf_counter() walltime_ms = (e2e_toc - e2e_tic) * 1000. if not warmup: self.print_summary(num_inference_steps, walltime_ms, batch_size) if not self.return_latents and save_image: # post-process images images = ((images + 1) * 255 / 2).clamp(0, 255).detach().permute(0, 2, 3, 1).round().type(torch.uint8).cpu().numpy() self.save_image(images, self.pipeline_type.name.lower(), prompt, self.seed) return (latents, walltime_ms) if self.return_latents else (images, walltime_ms) def run(self, prompt, negative_prompt, height, width, batch_size, batch_count, num_warmup_runs, use_cuda_graph, **kwargs): # 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): self.infer(prompt, negative_prompt, height, width, warmup=True, **kwargs) for _ in range(batch_count): print("[I] Running StableDiffusion pipeline") if self.nvtx_profile: cudart.cudaProfilerStart() self.infer(prompt, negative_prompt, height, width, warmup=False, **kwargs) if self.nvtx_profile: cudart.cudaProfilerStop()