# # 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 json import os import numpy as np import onnx import torch from diffusers import DiffusionPipeline from onnx import numpy_helper from demo_diffusion.model import load, optimizer from demo_diffusion.model.lora import merge_loras class BaseModel: def __init__( self, version="1.4", pipeline=None, device="cuda", hf_token="", verbose=True, framework_model_dir="pytorch_model", fp16=False, tf32=False, bf16=False, int8=False, fp8=False, fp4=False, max_batch_size=16, text_maxlen=77, embedding_dim=768, compression_factor=8, ): self.name = self.__class__.__name__ self.pipeline_type = pipeline self.pipeline = pipeline.name self.version = version self.path = load.get_path(version, pipeline) self.device = device self.hf_token = hf_token self.hf_safetensor = True self.verbose = verbose self.framework_model_dir = framework_model_dir self.fp16 = fp16 self.tf32 = tf32 self.bf16 = bf16 self.int8 = int8 self.fp8 = fp8 self.fp4 = fp4 self.compression_factor = compression_factor self.min_batch = 1 self.max_batch = max_batch_size self.min_image_shape = 256 # min image resolution: 256x256 self.max_image_shape = 1360 # max image resolution: 1360x1360 self.min_latent_shape = self.min_image_shape // self.compression_factor self.max_latent_shape = self.max_image_shape // self.compression_factor self.text_maxlen = text_maxlen self.embedding_dim = embedding_dim self.extra_output_names = [] self.do_constant_folding = True def get_pipeline(self): model_opts = {"variant": "fp16", "torch_dtype": torch.float16} if self.fp16 else {} model_opts = {"torch_dtype": torch.bfloat16} if self.bf16 else model_opts return DiffusionPipeline.from_pretrained( self.path, use_safetensors=self.hf_safetensor, token=self.hf_token, **model_opts, ).to(self.device) def get_model(self, torch_inference=""): pass def get_input_names(self): pass def get_output_names(self): pass def get_dynamic_axes(self): return None def get_sample_input(self, batch_size, image_height, image_width, static_shape): pass def get_input_profile(self, batch_size, image_height, image_width, static_batch, static_shape): return None def get_shape_dict(self, batch_size, image_height, image_width): return None # Helper utility for ONNX export def export_onnx( self, onnx_path, onnx_opt_path, onnx_opset, opt_image_height, opt_image_width, opt_num_frames=None, custom_model=None, enable_lora_merge=False, static_shape=False, lora_loader=None, dynamo=False, ): onnx_opt_graph = None # Export optimized ONNX model (if missing) if not os.path.exists(onnx_opt_path): if not os.path.exists(onnx_path): print(f"[I] Exporting ONNX model: {onnx_path}") def export_onnx(model): if enable_lora_merge: assert lora_loader is not None model = merge_loras(model, lora_loader) export_kwargs = {} if dynamo: export_kwargs["dynamic_shapes"] = self.get_dynamic_axes() else: export_kwargs["dynamic_axes"] = self.get_dynamic_axes() inputs = self.get_sample_input( 1, opt_image_height, opt_image_width, static_shape, **({'num_frames': opt_num_frames} if opt_num_frames else {}) ) with torch.no_grad(): torch.onnx.export( model, inputs, onnx_path, export_params=True, do_constant_folding=self.do_constant_folding, input_names=self.get_input_names(), output_names=self.get_output_names(), verbose=False, dynamo=dynamo, opset_version=onnx_opset, **export_kwargs, ) if custom_model: with torch.inference_mode(): export_onnx(custom_model) else: # WAR: Enable autocast for BF16 Stable Cascade pipeline do_autocast = True if self.version == "cascade" and self.bf16 else False model = self.get_model() with torch.inference_mode(), torch.autocast("cuda", enabled=do_autocast): export_onnx(model) del model gc.collect() torch.cuda.empty_cache() else: print(f"[I] Found cached ONNX model: {onnx_path}") print(f"[I] Optimizing ONNX model: {onnx_opt_path}") onnx_opt_graph = self.optimize(onnx.load(onnx_path)) if load.onnx_graph_needs_external_data(onnx_opt_graph): onnx.save_model( onnx_opt_graph, onnx_opt_path, save_as_external_data=True, all_tensors_to_one_file=True, convert_attribute=False, ) else: onnx.save(onnx_opt_graph, onnx_opt_path) else: print(f"[I] Found cached optimized ONNX model: {onnx_opt_path} ") # Helper utility for weights map def export_weights_map(self, onnx_opt_path, weights_map_path): if not os.path.exists(weights_map_path): onnx_opt_dir = os.path.dirname(onnx_opt_path) onnx_opt_model = onnx.load(onnx_opt_path) state_dict = self.get_model().state_dict() # Create initializer data hashes initializer_hash_mapping = {} for initializer in onnx_opt_model.graph.initializer: initializer_data = numpy_helper.to_array(initializer, base_dir=onnx_opt_dir).astype(np.float16) initializer_hash = hash(initializer_data.data.tobytes()) initializer_hash_mapping[initializer.name] = (initializer_hash, initializer_data.shape) weights_name_mapping = {} weights_shape_mapping = {} # set to keep track of initializers already added to the name_mapping dict initializers_mapped = set() for wt_name, wt in state_dict.items(): # get weight hash wt = wt.cpu().detach().numpy().astype(np.float16) wt_hash = hash(wt.data.tobytes()) wt_t_hash = hash(np.transpose(wt).data.tobytes()) for initializer_name, (initializer_hash, initializer_shape) in initializer_hash_mapping.items(): # Due to constant folding, some weights are transposed during export # To account for the transpose op, we compare the initializer hash to the # hash for the weight and its transpose if wt_hash == initializer_hash or wt_t_hash == initializer_hash: # The assert below ensures there is a 1:1 mapping between # PyTorch and ONNX weight names. It can be removed in cases where 1:many # mapping is found and name_mapping[wt_name] = list() assert initializer_name not in initializers_mapped weights_name_mapping[wt_name] = initializer_name initializers_mapped.add(initializer_name) is_transpose = False if wt_hash == initializer_hash else True weights_shape_mapping[wt_name] = (initializer_shape, is_transpose) # Sanity check: Were any weights not matched if wt_name not in weights_name_mapping: print(f"[I] PyTorch weight {wt_name} not matched with any ONNX initializer") print(f"[I] {len(weights_name_mapping.keys())} PyTorch weights were matched with ONNX initializers") assert weights_name_mapping.keys() == weights_shape_mapping.keys() with open(weights_map_path, "w") as fp: json.dump([weights_name_mapping, weights_shape_mapping], fp) else: print(f"[I] Found cached weights map: {weights_map_path} ") def optimize(self, onnx_graph, return_onnx=True, **kwargs): opt = optimizer.Optimizer(onnx_graph, verbose=self.verbose, version=self.version) opt.info(self.name + ": original") opt.cleanup() opt.info(self.name + ": cleanup") if kwargs.get("modify_fp8_graph", False): is_fp16_io = kwargs.get("is_fp16_io", True) opt.modify_fp8_graph(is_fp16_io=is_fp16_io) opt.info(self.name + ": modify fp8 graph") elif self.bf16: # Cast Resize I/O for strongly-typed TRT builds: BF16 -> FP32 inputs, FP32 -> BF16 outputs. # TRT does not support BF16 for the Resize operator. opt.infer_shapes() opt.cast_resize_io(output_dtype=onnx.TensorProto.BFLOAT16) opt.info(self.name + ": cast resize I/O for bf16") if self.version.startswith("flux.1") and self.fp8: opt.flux_convert_rope_weight_type() opt.info(self.name + ": convert rope weight type for fp8 flux") opt.fold_constants() opt.info(self.name + ": fold constants") opt.infer_shapes() opt.info(self.name + ": shape inference") if kwargs.get("modify_int8_graph", False): opt.modify_int8_graph() opt.info(self.name + ": modify int8 graph") onnx_opt_graph = opt.cleanup(return_onnx=return_onnx) opt.info(self.name + ": finished") return onnx_opt_graph def check_dims(self, batch_size, image_height, image_width, num_frames=None): assert batch_size >= self.min_batch and batch_size <= self.max_batch latent_height = image_height // self.compression_factor latent_width = image_width // self.compression_factor assert latent_height >= self.min_latent_shape and latent_height <= self.max_latent_shape assert latent_width >= self.min_latent_shape and latent_width <= self.max_latent_shape if num_frames: latent_frames = (self.num_frames - 1) // self.temporal_compression_factor + 1 return (latent_height, latent_width, latent_frames) return (latent_height, latent_width) def get_minmax_dims(self, batch_size, image_height, image_width, static_batch, static_shape, num_frames=None): min_batch = batch_size if static_batch else self.min_batch max_batch = batch_size if static_batch else self.max_batch latent_height = image_height // self.compression_factor latent_width = image_width // self.compression_factor min_image_height = image_height if static_shape else self.min_image_shape max_image_height = image_height if static_shape else self.max_image_shape min_image_width = image_width if static_shape else self.min_image_shape max_image_width = image_width if static_shape else self.max_image_shape min_latent_height = latent_height if static_shape else self.min_latent_shape max_latent_height = latent_height if static_shape else self.max_latent_shape min_latent_width = latent_width if static_shape else self.min_latent_shape max_latent_width = latent_width if static_shape else self.max_latent_shape frame_dims = () if num_frames: latent_frames = (num_frames - 1) // self.temporal_compression_factor + 1 min_latent_frames = latent_frames if static_shape else self.min_latent_frames max_latent_frames = latent_frames if static_shape else self.max_latent_frames frame_dims = (min_latent_frames, max_latent_frames) return ( min_batch, max_batch, min_image_height, max_image_height, min_image_width, max_image_width, min_latent_height, max_latent_height, min_latent_width, max_latent_width, *frame_dims )