wehub-resource-sync
321 lines
13 KiB
Python
321 lines
13 KiB
Python
#
|
|
# 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
|
|
)
|