# # 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. # from __future__ import annotations import abc import argparse import gc import json import os import pathlib import sys from abc import ABC, abstractmethod from typing import Any, List import modelopt.torch.opt as mto import modelopt.torch.quantization as mtq import nvtx import torch from cuda.bindings import runtime as cudart from diffusers import ( DDIMScheduler, DDPMScheduler, DDPMWuerstchenScheduler, EulerAncestralDiscreteScheduler, EulerDiscreteScheduler, FlowMatchEulerDiscreteScheduler, LCMScheduler, LMSDiscreteScheduler, PNDMScheduler, UniPCMultistepScheduler, ) from torch.utils.data import DataLoader import demo_diffusion.engine as engine_module import demo_diffusion.image as image_module from demo_diffusion.model import ( make_scheduler, merge_loras, unload_torch_model, ) from demo_diffusion.pipeline.calibrate import load_calib_prompts from demo_diffusion.pipeline.model_memory_manager import ModelMemoryManager from demo_diffusion.pipeline.type import PIPELINE_TYPE from demo_diffusion.utils_modelopt import ( SD_FP8_BF16_FLUX_MMDIT_BMM2_FP8_OUTPUT_CONFIG, SD_FP8_FP16_DEFAULT_CONFIG, SD_FP8_FP32_DEFAULT_CONFIG, PromptImageDataset, SameSizeSampler, check_lora, custom_collate, filter_func, filter_func_no_proj_out, fp8_mha_disable, generate_fp8_scales, get_int8_config, infinite_dataloader, quantize_lvl, set_fmha, set_quant_precision, ) class DiffusionPipeline(ABC): """ Application showcasing the acceleration of Stable Diffusion pipelines using NVidia TensorRT. """ VALID_DIFFUSION_PIPELINES = ( "1.4", "dreamshaper-7", "xl-1.0", "xl-turbo", "svd-xt-1.1", "sd3", "3.5-medium", "3.5-large", "cascade", "flux.1-dev", "flux.1-dev-canny", "flux.1-dev-depth", "flux.1-schnell", "flux.1-kontext-dev", "wan2.2-t2v-a14b", "cosmos-predict2-2b-text2image", "cosmos-predict2-14b-text2image", "cosmos-predict2-2b-video2world", "cosmos-predict2-14b-video2world", ) SCHEDULER_DEFAULTS = { "1.4": "PNDM", "dreamshaper-7": "PNDM", "xl-1.0": "Euler", "xl-turbo": "EulerA", "3.5-large": "FlowMatchEuler", "3.5-medium": "FlowMatchEuler", "svd-xt-1.1": "Euler", "cascade": "DDPMWuerstchen", "flux.1-dev": "FlowMatchEuler", "flux.1-dev-canny": "FlowMatchEuler", "flux.1-dev-depth": "FlowMatchEuler", "flux.1-schnell": "FlowMatchEuler", "flux.1-kontext-dev": "FlowMatchEuler", "wan2.2-t2v-a14b": "UniPC", "cosmos-predict2-2b-text2image": "FlowMatchEuler", "cosmos-predict2-14b-text2image": "FlowMatchEuler", "cosmos-predict2-2b-video2world": "FlowMatchEuler", "cosmos-predict2-14b-video2world": "FlowMatchEuler", } def __init__( self, dd_path, version="1.4", pipeline_type=PIPELINE_TYPE.TXT2IMG, bf16=False, max_batch_size=16, denoising_steps=30, scheduler=None, device="cuda", output_dir=".", hf_token=None, verbose=False, nvtx_profile=False, use_cuda_graph=False, framework_model_dir="pytorch_model", return_latents=False, low_vram=False, torch_inference="", torch_fallback=None, weight_streaming=False, text_encoder_weight_streaming_budget_percentage=None, denoiser_weight_streaming_budget_percentage=None, controlnet=None, ): """ Initializes the Diffusion pipeline. Args: dd_path (load_module.DDPath): DDPath object that contains all paths used in DemoDiffusion. version (str): The version of the pipeline. Should be one of the values listed in DiffusionPipeline.VALID_DIFFUSION_PIPELINES. pipeline_type (PIPELINE_TYPE): Task performed by the current pipeline. Should be one of PIPELINE_TYPE.__members__. max_batch_size (int): Maximum batch size for dynamic batch engine. bf16 (`bool`, defaults to False): Whether to run the pipeline in BFloat16 precision. 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 the values listed in DiffusionPipeline.SCHEDULER_DEFAULTS.values(). 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. framework_model_dir (str): cache directory for framework checkpoints. return_latents (bool): Skip decoding the image and return latents instead. low_vram (bool): [FLUX only] Optimize for low VRAM usage, possibly at the expense of inference performance. Disabled by default. torch_inference (str): Run inference with PyTorch (using specified compilation mode) instead of TensorRT. The compilation mode specified should be one of ['eager', 'reduce-overhead', 'max-autotune']. torch_fallback (str): [FLUX only] Comma separated list of models to be inferenced using PyTorch instead of TRT. For example --torch-fallback t5,transformer. If --torch-inference set, this parameter will be ignored. weight_streaming (`bool`, defaults to False): Whether to enable weight streaming during TensorRT engine build. text_encoder_ws_budget_percentage (`int`, defaults to None): Weight streaming budget as a percentage of the size of total streamable weights for the text encoder model. denoiser_weight_streaming_budget_percentage (`int`, defaults to None): Weight streaming budget as a percentage of the size of total streamable weights for the denoiser model. controlnet (str, defaults to None): Type of ControlNet to use for the pipeline. """ self.bf16 = bf16 self.dd_path = dd_path self.denoising_steps = denoising_steps 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.pipeline_type = pipeline_type self.return_latents = return_latents self.low_vram = low_vram self.weight_streaming = weight_streaming self.text_encoder_weight_streaming_budget_percentage = text_encoder_weight_streaming_budget_percentage self.denoiser_weight_streaming_budget_percentage = denoiser_weight_streaming_budget_percentage self.stages = self.get_model_names(self.pipeline_type, controlnet) # config to store additional info self.config = {} if torch_fallback: assert type(torch_fallback) is list for model_name in torch_fallback: if model_name not in self.stages: raise ValueError(f'Model "{model_name}" set in --torch-fallback does not exist') self.config[model_name.replace("-", "_") + "_torch_fallback"] = True print(f"[I] Setting torch_fallback for {model_name} model.") if not scheduler: scheduler = 'UniPC' if self.pipeline_type.is_controlnet() and not self.version == "3.5-large" 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, "FlowMatchEuler": FlowMatchEulerDiscreteScheduler, } try: scheduler_class = scheduler_class_map[scheduler] except KeyError: raise ValueError( f"Unsupported scheduler {scheduler}. Should be one of {list(scheduler_class_map.keys())}." ) self.scheduler = make_scheduler(scheduler_class, version, pipeline_type, hf_token, framework_model_dir) 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 load_engines() self.models = {} self.torch_models = {} self.engine = {} self.shape_dicts = {} self.shared_device_memory = None self.lora_loader = None # initialized in load_resources() self.events = {} self.generator = None self.markers = {} self.seed = None self.stream = None self.tokenizer = None def model_memory_manager(self, model_names, low_vram=False): return ModelMemoryManager(self, model_names, low_vram) @classmethod @abc.abstractmethod def FromArgs(cls, args: argparse.Namespace, pipeline_type: PIPELINE_TYPE) -> DiffusionPipeline: """Factory method to construct a concrete pipeline object from parsed arguments.""" raise NotImplementedError("FromArgs cannot be called from the abstract base class.") @classmethod @abc.abstractmethod def get_model_names(cls, pipeline_type: PIPELINE_TYPE, controlnet_type: str = None) -> List[str]: """Return a list of model names used by this pipeline.""" raise NotImplementedError("get_model_names cannot be called from the abstract base class.") @classmethod def _get_pipeline_uid(cls, version: str) -> str: """Return the unique ID of this pipeline. This is typically used to determine the default path for things like engine files, artifacts caches, etc. """ return f"{cls.__name__}_{version}" def profile_start(self, name, color="blue", domain=None): if self.nvtx_profile: self.markers[name] = nvtx.start_range(message=name, color=color, domain=domain) 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 load_resources(self, image_height, image_width, batch_size, seed): # Initialize noise generator if seed is not None: 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(): if self.torch_fallback[model_name]: continue shape_dict = obj.get_shape_dict( batch_size, image_height, image_width, **({'num_frames': obj.num_frames} if hasattr(obj, 'num_frames') and obj.num_frames else {}) ) self.shape_dicts[model_name] = shape_dict if not self.low_vram: self.engine[model_name].allocate_buffers(shape_dict=shape_dict, device=self.device) @abstractmethod def _initialize_models(self, *args, **kwargs): raise NotImplementedError("Please Implement the _initialize_models method") def _prepare_model_configs( self, enable_refit, int8, fp8, fp4 ): model_names = self.models.keys() self.torch_fallback = dict(zip(model_names, [self.torch_inference or self.config.get(model_name.replace('-','_')+'_torch_fallback', False) for model_name in model_names])) configs = {} for model_name in model_names: # Initialize config do_engine_refit = enable_refit and not self.pipeline_type.is_sd_xl_refiner() and any(model_name.startswith(prefix) for prefix in ("unet", "transformer")) do_lora_merge = not enable_refit and self.lora_loader and any(model_name.startswith(prefix) for prefix in ("unet", "transformer")) config = { "do_engine_refit": do_engine_refit, "do_lora_merge": do_lora_merge, "use_int8": False, "use_fp8": False, 'use_fp4': False, } # TODO: Move this to when arguments are first being validated in dd_argparse.py # 8-bit/4-bit precision inference if int8: assert self.pipeline_type.is_sd_xl_base() or self.version in [ "1.4", ], "int8 quantization only supported for SDXL and SD1.4 pipeline" if (self.pipeline_type.is_sd_xl() and model_name == "unetxl") or (model_name == "unet"): config["use_int8"] = True elif fp8: assert ( self.pipeline_type.is_sd_xl() or self.version in ["1.4"] or self.version.startswith("flux.1") or self.version.startswith("3.5-large") ), "fp8 quantization only supported for SDXL, SD1.4, SD3.5-large and FLUX pipelines" if ( (self.pipeline_type.is_sd_xl() and model_name == "unetxl") or (self.version.startswith("flux.1") and model_name == "transformer") or ( self.version.startswith("3.5-large") and ("transformer" in model_name or "controlnet" in model_name) ) or (model_name == "unet") ): config["use_fp8"] = True elif fp4: config['use_fp4'] = True # Setup paths config["onnx_path"] = self.dd_path.model_name_to_unoptimized_onnx_path[model_name] config["onnx_opt_path"] = self.dd_path.model_name_to_optimized_onnx_path[model_name] config["engine_path"] = self.dd_path.model_name_to_engine_path[model_name] config["weights_map_path"] = ( self.dd_path.model_name_to_weights_map_path[model_name] if config["do_engine_refit"] else None ) config["state_dict_path"] = self.dd_path.model_name_to_quantized_model_state_dict_path[model_name] config["refit_weights_path"] = self.dd_path.model_name_to_refit_weights_path[model_name] configs[model_name] = config return configs def _calibrate_and_save_model( self, pipeline, model, model_config, quantization_level, quantization_percentile, quantization_alpha, calibration_size, calib_batch_size, enable_lora_merge = False, **kwargs): print(f"[I] Calibrated weights not found, generating {model_config['state_dict_path']}") # 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 if kwargs["model_id"] in ("flux.1-dev", "flux.1-schnell"): common_args = { "prompt": prompts, "prompt_2": prompts, "num_inference_steps": kwargs["n_steps"], "height": kwargs.get("height", 1024), "width": kwargs.get("width", 1024), "guidance_scale": 3.5, "max_sequence_length": 512 if kwargs["model_id"] == "flux.1-dev" else 256, } else: common_args = { "prompt": prompts, "num_inference_steps": kwargs["n_steps"], "negative_prompt": ["normal quality, low quality, worst quality, low res, blurry, nsfw, nude"] * len(prompts), } pipeline(**common_args).images def do_calibrate_img2img(pipeline, dataloader, **kwargs): for i_th, (img_conds, prompts) in enumerate(dataloader): if i_th >= kwargs["calib_size"]: return common_args = { "prompt": list(prompts), "control_image": img_conds, "num_inference_steps": kwargs["n_steps"], "height": img_conds.size(2), "width": img_conds.size(3), "generator": torch.Generator().manual_seed(42), "guidance_scale": 3.5, "max_sequence_length": 512, } pipeline(**common_args).images if self.version in ("flux.1-dev-depth", "flux.1-dev-canny"): dataset = PromptImageDataset( root_dir=self.calibration_dataset, ) dataloader = DataLoader( dataset, batch_size=calib_batch_size, shuffle=False, num_workers=0, sampler=SameSizeSampler(dataset=dataset, batch_size=calib_batch_size), collate_fn=custom_collate, ) else: 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) def forward_loop(model): if self.version not in ("sd3", "flux.1-dev", "flux.1-schnell", "flux.1-dev-depth", "flux.1-dev-canny"): pipeline.unet = model else: pipeline.transformer = model if self.version in ("flux.1-dev-depth", "flux.1-dev-canny"): do_calibrate_img2img( pipeline=pipeline, dataloader=infinite_dataloader(dataloader), calib_size=calibration_size // calib_batch_size, n_steps=self.denoising_steps, model_id=self.version, ) else: do_calibrate( pipeline=pipeline, calibration_prompts=calibration_prompts, calib_size=calibration_size // calib_batch_size, n_steps=self.denoising_steps, model_id=self.version, **kwargs ) print(f"[I] Performing calibration for {calibration_size} steps.") if model_config['use_int8']: quant_config = get_int8_config( model, quantization_level, quantization_alpha, quantization_percentile, self.denoising_steps ) elif model_config['use_fp8']: if self.version.startswith("flux.1"): quant_config = SD_FP8_BF16_FLUX_MMDIT_BMM2_FP8_OUTPUT_CONFIG else: quant_config = SD_FP8_FP16_DEFAULT_CONFIG # Handle LoRA if enable_lora_merge: assert self.lora_loader is not None model = merge_loras(model, self.lora_loader) check_lora(model) if self.version.startswith("flux.1"): set_quant_precision(quant_config, "BFloat16") mtq.quantize(model, quant_config, forward_loop) mto.save(model, model_config['state_dict_path']) def _get_quantized_model( self, obj, model_config, quantization_level, quantization_percentile, quantization_alpha, calibration_size, calib_batch_size, enable_lora_merge = False, **kwargs): pipeline = obj.get_pipeline() is_flux = self.version.startswith("flux.1") model = pipeline.unet if self.version not in ("sd3", "flux.1-dev", "flux.1-schnell", "flux.1-dev-depth", "flux.1-dev-canny") else pipeline.transformer if model_config['use_fp8'] and quantization_level == 4.0: set_fmha(model, is_flux=is_flux) if not os.path.exists(model_config['state_dict_path']): self._calibrate_and_save_model( pipeline, model, model_config, quantization_level, quantization_percentile, quantization_alpha, calibration_size, calib_batch_size, enable_lora_merge, **kwargs) else: mto.restore(model, model_config['state_dict_path']) if not os.path.exists(model_config['onnx_path']): quantize_lvl(self.version, model, quantization_level) if self.version.startswith("flux.1"): mtq.disable_quantizer(model, filter_func_no_proj_out) else: mtq.disable_quantizer(model, filter_func) if model_config['use_fp8'] and not self.version.startswith("flux.1"): generate_fp8_scales(model) if quantization_level == 4.0: fp8_mha_disable(model, quantized_mha_output=False) # Remove Q/DQ after BMM2 in MHA else: model = None return model @abstractmethod def download_onnx_models(self, model_name: str, model_config: dict[str, Any]) -> None: """Download pre-exported ONNX Models""" raise NotImplementedError("Please Implement the download_onnx_models method") def is_native_export_supported(self, model_config: dict[str, Any]) -> bool: """Check if pipeline supports native ONNX export""" # Native export is supported by default return True @staticmethod def _fix_bf16_resize_nodes(onnx_opt_path): """Cast Resize node I/O for strongly-typed TRT engine builds. TRT does not support BF16 for the Resize operator, so inputs are cast to FP32 and outputs are cast back to BF16.""" import onnx import onnx_graphsurgeon as gs from demo_diffusion.model import load from demo_diffusion.utils_modelopt import cast_resize_io onnx_graph = onnx.load(onnx_opt_path, load_external_data=True) graph = gs.import_onnx(onnx_graph) resize_nodes = [n for n in graph.nodes if n.op == "Resize"] if not resize_nodes: return print(f"[I] Fixing {len(resize_nodes)} BF16 Resize node(s) in downloaded model: {onnx_opt_path}") cast_resize_io(graph, output_dtype=onnx.TensorProto.BFLOAT16) graph.cleanup().toposort() onnx_graph = gs.export_onnx(graph) if load.onnx_graph_needs_external_data(onnx_graph): onnx.save_model( onnx_graph, onnx_opt_path, save_as_external_data=True, all_tensors_to_one_file=True, convert_attribute=False, ) else: onnx.save(onnx_graph, onnx_opt_path) def _export_onnx( self, obj, model_name, model_config, opt_image_height, opt_image_width, static_shape, onnx_opset, quantization_level, quantization_percentile, quantization_alpha, calibration_size, calib_batch_size, onnx_export_only, download_onnx_models, ): # With onnx_export_only True, the export still happens even if the TRT engine exists. However, it will not re-run the export if the onnx exists. do_export_onnx = (not os.path.exists(model_config['engine_path']) or onnx_export_only) and not os.path.exists(model_config['onnx_opt_path']) do_export_weights_map = model_config['weights_map_path'] and not os.path.exists(model_config['weights_map_path']) # If ONNX export is required, either download ONNX models or check if the pipeline supports native ONNX export if do_export_onnx: if download_onnx_models: self.download_onnx_models(model_name, model_config) # Fix Resize nodes for strongly-typed TRT builds. # Downloaded models bypass optimize(), so apply the fix here. if obj.bf16: self._fix_bf16_resize_nodes(model_config['onnx_opt_path']) do_export_onnx = False else: self.is_native_export_supported(model_config) dynamo = True if (self.pipeline_type.is_video2world() and model_name == "transformer") or (self.pipeline_type.is_txt2vid() and (model_name in ["transformer", "transformer_2"])) or (self.version.startswith("flux.1") and model_name == "transformer" and obj.fp16) else False export_kwargs = { "static_shape": static_shape, "dynamo": dynamo, **({'opt_num_frames': obj.num_frames} if hasattr(obj, 'num_frames') and obj.num_frames else {}) } if do_export_onnx or do_export_weights_map: if not model_config['use_int8'] and not model_config['use_fp8']: obj.export_onnx( model_config["onnx_path"], model_config["onnx_opt_path"], onnx_opset, opt_image_height, opt_image_width, enable_lora_merge=model_config["do_lora_merge"], lora_loader=self.lora_loader, **export_kwargs, ) else: print(f"[I] Generating quantized ONNX model: {model_config['onnx_path']}") quantized_model = self._get_quantized_model( obj, model_config, quantization_level, quantization_percentile, quantization_alpha, calibration_size, calib_batch_size, height=opt_image_width, width=opt_image_width, enable_lora_merge=model_config["do_lora_merge"], ) obj.export_onnx( model_config["onnx_path"], model_config["onnx_opt_path"], onnx_opset, opt_image_height, opt_image_width, custom_model=quantized_model, **export_kwargs, ) # FIXME do_export_weights_map needs ONNX graph if do_export_weights_map: print(f"[I] Saving weights map: {model_config['weights_map_path']}") obj.export_weights_map(model_config['onnx_opt_path'], model_config['weights_map_path']) def _build_engine(self, obj, engine, model_config, opt_batch_size, opt_image_height, opt_image_width, optimization_level, static_batch, static_shape, enable_all_tactics, timing_cache): update_output_names = obj.get_output_names() + obj.extra_output_names if obj.extra_output_names else None tf32amp = obj.tf32 weight_streaming = getattr(obj, 'weight_streaming', False) precision_constraints = 'none' input_profile = obj.get_input_profile( opt_batch_size, opt_image_height, opt_image_width, static_batch=static_batch, static_shape=static_shape, **({'num_frames': obj.num_frames} if hasattr(obj, 'num_frames') and obj.num_frames else {}) ) engine.build( model_config["onnx_opt_path"], tf32=tf32amp, input_profile=input_profile, enable_refit=model_config["do_engine_refit"], enable_all_tactics=enable_all_tactics, timing_cache=timing_cache, update_output_names=update_output_names, weight_streaming=weight_streaming, verbose=self.verbose, builder_optimization_level=optimization_level, precision_constraints=precision_constraints, ) def _refit_engine(self, obj, model_name, model_config): assert model_config['weights_map_path'] with open(model_config['weights_map_path'], 'r') as fp_wts: print(f"[I] Loading weights map: {model_config['weights_map_path']} ") [weights_name_mapping, weights_shape_mapping] = json.load(fp_wts) if not os.path.exists(model_config['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(), model_config["onnx_opt_path"], weights_name_mapping, weights_shape_mapping ) print(f"[I] Saving refit weights: {model_config['refit_weights_path']}") torch.save((refit_weights, updated_weight_names), model_config["refit_weights_path"]) unload_torch_model(model) else: print(f"[I] Loading refit weights: {model_config['refit_weights_path']}") refit_weights, updated_weight_names = torch.load(model_config['refit_weights_path']) self.engine[model_name].refit(refit_weights, updated_weight_names) def _load_torch_models(self): # Load torch models for model_name, obj in self.models.items(): if self.torch_fallback[model_name]: self.torch_models[model_name] = obj.get_model(torch_inference=self.torch_inference) if self.low_vram: self.torch_models[model_name] = self.torch_models[model_name].to('cpu') torch.cuda.empty_cache() def load_engines( self, framework_model_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, fp4=False, quantization_level=2.5, quantization_percentile=1.0, quantization_alpha=0.8, calibration_size=32, calib_batch_size=2, onnx_export_only=False, download_onnx_models=False, ): """ Build and load engines for TensorRT accelerated inference. Export ONNX models first, if applicable. Args: framework_model_dir (str): Directory to store the framework model ckpt. 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. onnx_export_only (bool): Whether only export onnx without building the TRT engine. download_onnx_models (bool): Download pre-exported ONNX models """ self._initialize_models(framework_model_dir, int8, fp8, fp4) model_configs = self._prepare_model_configs(enable_refit, int8, fp8, fp4) # Export models to ONNX for model_name, obj in self.models.items(): if self.torch_fallback[model_name]: continue self._export_onnx( obj, model_name, model_configs[model_name], opt_image_height, opt_image_width, static_shape, onnx_opset, quantization_level, quantization_percentile, quantization_alpha, calibration_size, calib_batch_size, onnx_export_only, download_onnx_models, ) # Release temp GPU memory during onnx export to avoid OOM. gc.collect() torch.cuda.empty_cache() if onnx_export_only: return # Build TensorRT engines for model_name, obj in self.models.items(): if self.torch_fallback[model_name]: continue model_config = model_configs[model_name] engine = engine_module.Engine(model_config["engine_path"]) if not os.path.exists(model_config['engine_path']): self._build_engine(obj, engine, model_config, opt_batch_size, opt_image_height, opt_image_width, optimization_level, static_batch, static_shape, enable_all_tactics, timing_cache) self.engine[model_name] = engine # Load and refit TensorRT engines for model_name, obj in self.models.items(): if self.torch_fallback[model_name]: continue model_config = model_configs[model_name] # For non low_vram case, the engines will remain in GPU memory from now on. assert self.engine[model_name].engine is None if not self.low_vram: weight_streaming = getattr(obj, 'weight_streaming', False) weight_streaming_budget_percentage = getattr(obj, 'weight_streaming_budget_percentage', None) self.engine[model_name].load(weight_streaming, weight_streaming_budget_percentage) if model_config['do_engine_refit'] and self.lora_loader: # For low_vram, using on-demand load and unload for refit. if self.low_vram: assert self.engine[model_name].engine is None self.engine[model_name].load() self._refit_engine(obj, model_name, model_config) if self.low_vram: self.engine[model_name].unload() # Load PyTorch models if torch-inference mode is enabled self._load_torch_models() # Reclaim GPU memory from torch cache torch.cuda.empty_cache() def calculate_max_device_memory(self): max_device_memory = 0 for model_name, engine in self.engine.items(): if self.low_vram: engine.load() max_device_memory = max(max_device_memory, engine.engine.device_memory_size_v2) if self.low_vram: engine.unload() return max_device_memory def get_device_memory_sizes(self): device_memory_sizes = {} for model_name, engine in self.engine.items(): engine.load() device_memory_sizes[model_name] = engine.engine.device_memory_size_v2 engine.unload() return device_memory_sizes def activate_engines(self, shared_device_memory=None): if shared_device_memory is None: max_device_memory = self.calculate_max_device_memory() _, shared_device_memory = cudart.cudaMalloc(max_device_memory) self.shared_device_memory = shared_device_memory # Load and activate TensorRT engines if not self.low_vram: for engine in self.engine.values(): engine.activate(device_memory=self.shared_device_memory) def run_engine(self, model_name, feed_dict): engine = self.engine[model_name] # CUDA graphs should be disabled when low_vram is enabled. if self.low_vram: assert self.use_cuda_graph == False return engine.infer(feed_dict, self.stream, use_cuda_graph=self.use_cuda_graph) def teardown(self): for e in self.events.values(): cudart.cudaEventDestroy(e[0]) cudart.cudaEventDestroy(e[1]) for engine in self.engine.values(): engine.deallocate_buffers() engine.deactivate() engine.unload(verbose=False) del engine if self.shared_device_memory: cudart.cudaFree(self.shared_device_memory) for torch_model in self.torch_models.values(): torch_model.to("cpu") del torch_model cudart.cudaStreamDestroy(self.stream) del self.stream gc.collect() torch.cuda.empty_cache() def initialize_latents(self, batch_size, unet_channels, latent_height, latent_width, latents_dtype=torch.float32): 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 save_image(self, images, pipeline, prompt, seed): # Save image prompt_prefix = ''.join(set([prompt[i].replace(' ','_')[:10] for i in range(len(prompt))])) image_name_prefix = '-'.join([pipeline, prompt_prefix, 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) @abstractmethod def print_summary(self): """Print a summary of the pipeline's configuration.""" raise NotImplementedError("Please Implement the print_summary method") @abstractmethod def infer(self): """Perform inference using the pipeline.""" raise NotImplementedError("Please Implement the infer method") @abstractmethod def run(self): """Run the pipeline.""" raise NotImplementedError("Please Implement the run method")