mirrors/nvidia--tensorrt
1
0
Fork 0
Code Issues Pull Requests Actions 9 Deployments Models Agent Notes Packages Projects Releases Wiki Activity

chore: import upstream snapshot with attribution
Docker Image CI / build-ubuntu2004 (push) Has been cancelled
Details

Browse Source
This commit is contained in:
wehub-resource-sync
2026-07-13 13:36:55 +08:00
commit c8a779b1bb
1887 changed files with 3245738 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
demo/Diffusion/demo_diffusion/utils_modelopt.py
Executable
+853
View File
@@ -0,0 +1,853 @@
#
# 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 os
import re
from collections import defaultdict
from random import choice, shuffle
from typing import Set
import modelopt.torch.quantization as mtq
import numpy as np
import onnx_graphsurgeon as gs
import torch
import torch.nn.functional as F
from diffusers.models.attention_processor import (
Attention,
AttnProcessor,
FluxAttnProcessor2_0,
)
from diffusers.models.lora import LoRACompatibleConv, LoRACompatibleLinear
from modelopt.torch.quantization import utils as quant_utils
from modelopt.torch.quantization.calib.max import MaxCalibrator
from PIL import Image
from torch.utils.data import Dataset, Sampler
import onnx
USE_PEFT = True
try:
from peft.tuners.lora.layer import Conv2d as PEFTLoRAConv2d
from peft.tuners.lora.layer import Linear as PEFTLoRALinear
except ModuleNotFoundError:
USE_PEFT = False
class PercentileCalibrator(MaxCalibrator):
def __init__(self, num_bits=8, axis=None, unsigned=False, track_amax=False, **kwargs):
super().__init__(num_bits, axis, unsigned, track_amax)
self.percentile = kwargs["percentile"]
self.total_step = kwargs["total_step"]
self.collect_method = kwargs["collect_method"]
self.data = {}
self.i = 0
def collect(self, x):
"""Tracks the absolute max of all tensors.
Args:
x: A tensor
Raises:
RuntimeError: If amax shape changes
"""
# Swap axis to reduce.
axis = self._axis if isinstance(self._axis, (list, tuple)) else [self._axis]
# Handle negative axis.
axis = [x.dim() + i if isinstance(i, int) and i < 0 else i for i in axis]
reduce_axis = []
for i in range(x.dim()):
if i not in axis:
reduce_axis.append(i)
local_amax = quant_utils.reduce_amax(x, axis=reduce_axis).detach()
_cur_step = self.i % self.total_step
if _cur_step not in self.data.keys():
self.data[_cur_step] = local_amax
else:
if self.collect_method == "global_min":
self.data[_cur_step] = torch.min(self.data[_cur_step], local_amax)
elif self.collect_method == "min-max" or self.collect_method == "mean-max":
self.data[_cur_step] = torch.max(self.data[_cur_step], local_amax)
else:
self.data[_cur_step] += local_amax
if self._track_amax:
raise NotImplementedError
self.i += 1
def compute_amax(self):
"""Return the absolute max of all tensors collected."""
up_lim = int(self.total_step * self.percentile)
if self.collect_method == "min-mean":
amaxs_values = [self.data[i] / self.total_step for i in range(0, up_lim)]
else:
amaxs_values = [self.data[i] for i in range(0, up_lim)]
if self.collect_method == "mean-max":
act_amax = torch.vstack(amaxs_values).mean(axis=0)[0]
else:
act_amax = torch.vstack(amaxs_values).min(axis=0)[0]
self._calib_amax = act_amax
return self._calib_amax
def __str__(self):
s = "PercentileCalibrator"
return s.format(**self.__dict__)
def __repr__(self):
s = "PercentileCalibrator("
s += super(MaxCalibrator, self).__repr__()
s += " calib_amax={_calib_amax}"
if self._track_amax:
s += " amaxs={_amaxs}"
s += ")"
return s.format(**self.__dict__)
def filter_func(name):
pattern = re.compile(
r".*(time_emb_proj|time_embedding|conv_in|conv_out|conv_shortcut|add_embedding|pos_embed|time_text_embed|context_embedder|norm_out|proj_out).*"
)
return pattern.match(name) is not None
def filter_func_no_proj_out(name): # used for Flux
pattern = re.compile(
r".*(time_emb_proj|time_embedding|conv_in|conv_out|conv_shortcut|add_embedding|pos_embed|time_text_embed|context_embedder|norm_out|x_embedder).*"
)
return pattern.match(name) is not None
def quantize_lvl(model_id, backbone, quant_level=2.5, linear_only=False, enable_conv_3d=True):
"""
We should disable the unwanted quantizer when exporting the onnx
Because in the current modelopt setting, it will load the quantizer amax for all the layers even
if we didn't add that unwanted layer into the config during the calibration
"""
for name, module in backbone.named_modules():
if isinstance(module, torch.nn.Conv2d):
if linear_only:
module.input_quantizer.disable()
module.weight_quantizer.disable()
else:
module.input_quantizer.enable()
module.weight_quantizer.enable()
elif isinstance(module, torch.nn.Linear):
if (
(quant_level >= 2 and "ff.net" in name)
or (quant_level >= 2.5 and ("to_q" in name or "to_k" in name or "to_v" in name))
or quant_level >= 3
) and name != "proj_out": # Disable the final output layer from flux model
module.input_quantizer.enable()
module.weight_quantizer.enable()
else:
module.input_quantizer.disable()
module.weight_quantizer.disable()
elif isinstance(module, torch.nn.Conv3d) and not enable_conv_3d:
"""
Error: Torch bug, ONNX export failed due to unknown kernel shape in QuantConv3d.
TRT_FP8QuantizeLinear and TRT_FP8DequantizeLinear operations in UNetSpatioTemporalConditionModel for svd
cause issues. Inputs on different devices (CUDA vs CPU) may contribute to the problem.
"""
module.input_quantizer.disable()
module.weight_quantizer.disable()
elif isinstance(module, Attention):
# TRT only supports FP8 MHA with head_size % 16 == 0.
head_size = int(module.inner_dim / module.heads)
if quant_level >= 4 and head_size % 16 == 0:
module.q_bmm_quantizer.enable()
module.k_bmm_quantizer.enable()
module.v_bmm_quantizer.enable()
module.softmax_quantizer.enable()
if model_id.startswith("flux.1"):
if name.startswith("transformer_blocks"):
module.bmm2_output_quantizer.enable()
else:
module.bmm2_output_quantizer.disable()
setattr(module, "_disable_fp8_mha", False)
else:
module.q_bmm_quantizer.disable()
module.k_bmm_quantizer.disable()
module.v_bmm_quantizer.disable()
module.softmax_quantizer.disable()
module.bmm2_output_quantizer.disable()
setattr(module, "_disable_fp8_mha", True)
def fp8_mha_disable(backbone, quantized_mha_output: bool = True):
def mha_filter_func(name):
pattern = re.compile(
r".*(q_bmm_quantizer|k_bmm_quantizer|v_bmm_quantizer|softmax_quantizer).*"
if quantized_mha_output
else r".*(q_bmm_quantizer|k_bmm_quantizer|v_bmm_quantizer|softmax_quantizer|bmm2_output_quantizer).*"
)
return pattern.match(name) is not None
if hasattr(F, "scaled_dot_product_attention"):
mtq.disable_quantizer(backbone, mha_filter_func)
def get_int8_config(
model,
quant_level=3,
alpha=0.8,
percentile=1.0,
num_inference_steps=20,
collect_method="min-mean",
):
quant_config = {
"quant_cfg": {
"*lm_head*": {"enable": False},
"*output_layer*": {"enable": False},
"*output_quantizer": {"enable": False},
"default": {"num_bits": 8, "axis": None},
},
"algorithm": {"method": "smoothquant", "alpha": alpha},
}
for name, module in model.named_modules():
w_name = f"{name}*weight_quantizer"
i_name = f"{name}*input_quantizer"
if w_name in quant_config["quant_cfg"].keys() or i_name in quant_config["quant_cfg"].keys():
continue
if filter_func(name):
continue
if isinstance(module, (torch.nn.Linear, LoRACompatibleLinear)):
if (
(quant_level >= 2 and "ff.net" in name)
or (quant_level >= 2.5 and ("to_q" in name or "to_k" in name or "to_v" in name))
or quant_level == 3
):
quant_config["quant_cfg"][w_name] = {"num_bits": 8, "axis": 0}
quant_config["quant_cfg"][i_name] = {"num_bits": 8, "axis": -1}
elif isinstance(module, (torch.nn.Conv2d, LoRACompatibleConv)):
quant_config["quant_cfg"][w_name] = {"num_bits": 8, "axis": 0}
quant_config["quant_cfg"][i_name] = {
"num_bits": 8,
"axis": None,
"calibrator": (
PercentileCalibrator,
(),
{
"num_bits": 8,
"axis": None,
"percentile": percentile,
"total_step": num_inference_steps,
"collect_method": collect_method,
},
),
}
return quant_config
SD_FP8_FP16_DEFAULT_CONFIG = {
"quant_cfg": {
"*weight_quantizer": {"num_bits": (4, 3), "axis": None, "trt_high_precision_dtype": "Half"},
"*input_quantizer": {"num_bits": (4, 3), "axis": None, "trt_high_precision_dtype": "Half"},
"*output_quantizer": {"enable": False},
"*q_bmm_quantizer": {"num_bits": (4, 3), "axis": None, "trt_high_precision_dtype": "Half"},
"*k_bmm_quantizer": {"num_bits": (4, 3), "axis": None, "trt_high_precision_dtype": "Half"},
"*v_bmm_quantizer": {"num_bits": (4, 3), "axis": None, "trt_high_precision_dtype": "Half"},
"*softmax_quantizer": {
"num_bits": (4, 3),
"axis": None,
"trt_high_precision_dtype": "Half",
},
"default": {"enable": False},
},
"algorithm": "max",
}
SD_FP8_BF16_DEFAULT_CONFIG = {
"quant_cfg": {
"*weight_quantizer": {"num_bits": (4, 3), "axis": None, "trt_high_precision_dtype": "BFloat16"},
"*input_quantizer": {"num_bits": (4, 3), "axis": None, "trt_high_precision_dtype": "BFloat16"},
"*output_quantizer": {"enable": False},
"*q_bmm_quantizer": {"num_bits": (4, 3), "axis": None, "trt_high_precision_dtype": "BFloat16"},
"*k_bmm_quantizer": {"num_bits": (4, 3), "axis": None, "trt_high_precision_dtype": "BFloat16"},
"*v_bmm_quantizer": {"num_bits": (4, 3), "axis": None, "trt_high_precision_dtype": "BFloat16"},
"*softmax_quantizer": {
"num_bits": (4, 3),
"axis": None,
"trt_high_precision_dtype": "BFloat16",
},
"default": {"enable": False},
},
"algorithm": "max",
}
SD_FP8_BF16_FLUX_MMDIT_BMM2_FP8_OUTPUT_CONFIG = {
"quant_cfg": {
"*weight_quantizer": {"num_bits": (4, 3), "axis": None, "trt_high_precision_dtype": "BFloat16"},
"*input_quantizer": {"num_bits": (4, 3), "axis": None, "trt_high_precision_dtype": "BFloat16"},
"*output_quantizer": {"enable": False},
"*q_bmm_quantizer": {"num_bits": (4, 3), "axis": None, "trt_high_precision_dtype": "BFloat16"},
"*k_bmm_quantizer": {"num_bits": (4, 3), "axis": None, "trt_high_precision_dtype": "BFloat16"},
"*v_bmm_quantizer": {"num_bits": (4, 3), "axis": None, "trt_high_precision_dtype": "BFloat16"},
"*softmax_quantizer": {
"num_bits": (4, 3),
"axis": None,
"trt_high_precision_dtype": "BFloat16",
},
"transformer_blocks*bmm2_output_quantizer": {
"num_bits": (4, 3),
"axis": None,
"trt_high_precision_dtype": "BFloat16",
},
"default": {"enable": False},
},
"algorithm": "max",
}
SD_FP8_FP32_DEFAULT_CONFIG = {
"quant_cfg": {
"*weight_quantizer": {"num_bits": (4, 3), "axis": None, "trt_high_precision_dtype": "Float"},
"*input_quantizer": {"num_bits": (4, 3), "axis": None, "trt_high_precision_dtype": "Float"},
"*output_quantizer": {"enable": False},
"*q_bmm_quantizer": {"num_bits": (4, 3), "axis": None, "trt_high_precision_dtype": "Float"},
"*k_bmm_quantizer": {"num_bits": (4, 3), "axis": None, "trt_high_precision_dtype": "Float"},
"*v_bmm_quantizer": {"num_bits": (4, 3), "axis": None, "trt_high_precision_dtype": "Float"},
"*softmax_quantizer": {
"num_bits": (4, 3),
"axis": None,
"trt_high_precision_dtype": "Float",
},
"default": {"enable": False},
},
"algorithm": "max",
}
def set_fmha(denoiser, is_flux=False):
for name, module in denoiser.named_modules():
if isinstance(module, Attention):
if is_flux:
module.set_processor(FluxAttnProcessor2_0())
else:
module.set_processor(AttnProcessor())
def check_lora(model):
for name, module in model.named_modules():
if isinstance(module, (LoRACompatibleConv, LoRACompatibleLinear)):
assert (
module.lora_layer is None
), f"To quantize {name}, LoRA layer should be fused/merged. Please fuse the LoRA layer before quantization."
elif USE_PEFT and isinstance(module, (PEFTLoRAConv2d, PEFTLoRALinear)):
assert (
module.merged
), f"To quantize {name}, LoRA layer should be fused/merged. Please fuse the LoRA layer before quantization."
def generate_fp8_scales(unet):
# temporary solution due to a known bug in torch.onnx._dynamo_export
for _, module in unet.named_modules():
if isinstance(module, (torch.nn.Linear, torch.nn.Conv2d)) and (
hasattr(module.input_quantizer, "_amax") and module.input_quantizer is not None
):
module.input_quantizer._num_bits = 8
module.weight_quantizer._num_bits = 8
module.input_quantizer._amax = module.input_quantizer._amax * (127 / 448.0)
module.weight_quantizer._amax = module.weight_quantizer._amax * (127 / 448.0)
elif isinstance(module, Attention) and (
hasattr(module.q_bmm_quantizer, "_amax")
and module.q_bmm_quantizer is not None
and hasattr(module.k_bmm_quantizer, "_amax")
and module.k_bmm_quantizer is not None
and hasattr(module.v_bmm_quantizer, "_amax")
and module.v_bmm_quantizer is not None
and hasattr(module.softmax_quantizer, "_amax")
and module.softmax_quantizer is not None
):
module.q_bmm_quantizer._num_bits = 8
module.q_bmm_quantizer._amax = module.q_bmm_quantizer._amax * (127 / 448.0)
module.k_bmm_quantizer._num_bits = 8
module.k_bmm_quantizer._amax = module.k_bmm_quantizer._amax * (127 / 448.0)
module.v_bmm_quantizer._num_bits = 8
module.v_bmm_quantizer._amax = module.v_bmm_quantizer._amax * (127 / 448.0)
module.softmax_quantizer._num_bits = 8
module.softmax_quantizer._amax = module.softmax_quantizer._amax * (127 / 448.0)
def get_parent_nodes(node):
"""
Returns list of input producer nodes for the given node.
"""
parents = []
for tensor in node.inputs:
# If the tensor is not a constant or graph input and has a producer,
# the producer is a parent of node `node`
if len(tensor.inputs) == 1:
parents.append(tensor.inputs[0])
return parents
def get_child_nodes(node):
"""
Returns list of output consumer nodes for the given node.
"""
children = []
for tensor in node.outputs:
for consumer in tensor.outputs: # Traverse all consumer of the tensor
children.append(consumer)
return children
def has_path_type(node, graph, path_type, is_forward, wild_card_types, path_nodes):
"""
Return pattern nodes for the given path_type.
"""
if not path_type:
# All types matched
return True
# Check if current non-wild node type does not match the expected path type
node_type = node.op
is_match = node_type == path_type[0]
is_wild_match = node_type in wild_card_types
if not is_match and not is_wild_match:
return False
if is_match:
path_nodes.append(node)
next_path_type = path_type[1:]
else:
next_path_type = path_type[:]
if is_forward:
next_level_nodes = get_child_nodes(node)
else:
next_level_nodes = get_parent_nodes(node)
# Check if any child (forward path) or parent (backward path) can match the remaining path types
for next_node in next_level_nodes:
sub_path = []
if has_path_type(next_node, graph, next_path_type, is_forward, wild_card_types, sub_path):
path_nodes.extend(sub_path)
return True
# Path type matches if there is no remaining types to match
return not next_path_type
def insert_cast(graph, input_tensor, attrs):
"""
Create a cast layer using tensor as input.
"""
output_tensor = gs.Variable(name=f"{input_tensor.name}/Cast_output", dtype=attrs["to"])
next_node_list = input_tensor.outputs.copy()
graph.layer(
op="Cast",
name=f"{input_tensor.name}/Cast",
inputs=[input_tensor],
outputs=[output_tensor],
attrs=attrs,
)
# use cast output as input to next node
for next_node in next_node_list:
for idx, next_input in enumerate(next_node.inputs):
if next_input.name == input_tensor.name:
next_node.inputs[idx] = output_tensor
def cast_layernorm_io(graph):
"""
Cast LayerNormalization scale and bias inputs from FP16 to FP32.
In INT8 quantized graphs, DequantizeLinear outputs Float32 activations,
but LayerNorm scale/bias remain FP16 from the original model, causing
a type mismatch with --strongly-typed TensorRT builds.
"""
layernorm_nodes = [node for node in graph.nodes if node.op == "LayerNormalization"]
print(f"Found {len(layernorm_nodes)} LayerNormalization nodes to fix")
for node in layernorm_nodes:
# LayerNormalization inputs: 0=X (data), 1=Scale, 2=B (bias, optional)
for i in range(1, len(node.inputs)):
input_tensor = node.inputs[i]
if input_tensor.name and hasattr(input_tensor, 'dtype') and input_tensor.dtype == np.float16:
insert_cast(graph, input_tensor=input_tensor, attrs={"to": np.float32})
def cast_convtranspose_io(graph):
"""
Fix ConvTranspose input/output type mismatches for strongly-typed TRT builds.
In mixed-precision graphs (e.g. BF16 Stable Cascade VQGAN), architectural FP16->FP32
casts can leave a ConvTranspose with a FP32 activation input but FP16 kernel weights.
We cast the activation to match the kernel dtype, then cast the output back to the
original activation dtype so surrounding FP32 ops (e.g. residual Add) are unaffected.
"""
convtranspose_nodes = [node for node in graph.nodes if node.op == "ConvTranspose"]
fixed = 0
for node in convtranspose_nodes:
if len(node.inputs) < 2:
continue
act_input = node.inputs[0]
kernel = node.inputs[1]
if act_input.dtype is None or kernel.dtype is None or act_input.dtype == kernel.dtype:
continue
orig_dtype = act_input.dtype # e.g. np.dtype('float32')
target_dtype = kernel.dtype.type # e.g. np.float16
insert_cast(graph, input_tensor=act_input, attrs={"to": target_dtype})
# Update the output dtype to match and cast back, so downstream FP32 ops are unaffected.
for out in node.outputs:
if out.name and out.dtype == orig_dtype:
out.dtype = target_dtype
insert_cast(graph, input_tensor=out, attrs={"to": orig_dtype.type})
fixed += 1
print(f"Fixed {fixed} ConvTranspose input/output type mismatches")
def convert_zp_fp8(onnx_graph):
"""
Convert Q/DQ zero datatype from INT8 to FP8.
"""
# Find all zero constant nodes
qdq_zero_nodes = set()
for node in onnx_graph.graph.node:
if node.op_type == "QuantizeLinear":
if len(node.input) > 2:
qdq_zero_nodes.add(node.input[2])
print(f"Found {len(qdq_zero_nodes)} QDQ pairs")
# Convert zero point datatype from INT8 to FP8.
for node in onnx_graph.graph.node:
if node.output[0] in qdq_zero_nodes:
node.attribute[0].t.data_type = onnx.TensorProto.FLOAT8E4M3FN
return onnx_graph
def cast_resize_io(graph, output_dtype=np.float16):
"""
Add cast nodes to Resize nodes I/O because Resize needs to be run in fp32.
Inputs are cast to FP32, outputs are cast back to output_dtype (FP16 or BF16).
"""
resize_nodes = [node for node in graph.nodes if node.op == "Resize"]
print(f"Found {len(resize_nodes)} Resize nodes to fix")
for resize_node in resize_nodes:
# Skip Resize nodes whose data input is already FP32 — no casting needed.
if resize_node.inputs[0].dtype == np.float32:
continue
for i, input_tensor in enumerate(resize_node.inputs):
SIZES_INPUT_INDEX = 3 # Optional input "sizes" at index 3 must be in INT64. Skip cast for this input.
if i != SIZES_INPUT_INDEX and input_tensor.name:
insert_cast(graph, input_tensor=input_tensor, attrs={"to": np.float32})
for output_tensor in resize_node.outputs:
if output_tensor.name:
insert_cast(graph, input_tensor=output_tensor, attrs={"to": output_dtype})
def cast_fp8_mha_io(graph):
r"""
Insert three cast ops.
The first cast will be added before the input0 of MatMul to cast fp16 to fp32.
The second cast will be added before the input1 of MatMul to cast fp16 to fp32.
The third cast will be added after the output of MatMul to cast fp32 back to fp16.
Q Q
| |
DQ DQ
| |
Cast Cast
(fp16 to fp32) (fp16 to fp32)
\ /
\ /
\ /
MatMul
|
Cast (fp32 to fp16)
|
Q
|
DQ
The insertion of Cast ops in the FP8 MHA part actually forbids the MHAs to run
with FP16 accumulation because TensorRT only has FP32 accumulation kernels for FP8 MHAs.
"""
# Find FP8 MHA pattern.
# Match FP8 MHA: Q -> DQ -> BMM1 -> (Mul/Div) -> (Add) -> Softmax -> (Cast) -> Q -> DQ -> BMM2 -> Q -> DQ
softmax_bmm1_chain_type = ["Softmax", "MatMul", "DequantizeLinear", "QuantizeLinear"]
softmax_bmm2_chain_type = [
"Softmax",
"QuantizeLinear",
"DequantizeLinear",
"MatMul",
"QuantizeLinear",
"DequantizeLinear",
]
wild_card_types = [
"Div",
"Mul",
"ConstMul",
"Add",
"BiasAdd",
"Reshape",
"Transpose",
"Flatten",
"Cast",
]
fp8_mha_partitions = []
for node in graph.nodes:
if node.op == "Softmax":
fp8_mha_partition = []
if has_path_type(
node, graph, softmax_bmm1_chain_type, False, wild_card_types, fp8_mha_partition
) and has_path_type(
node, graph, softmax_bmm2_chain_type, True, wild_card_types, fp8_mha_partition
):
if (
len(fp8_mha_partition) == 10
and fp8_mha_partition[1].op == "MatMul"
and fp8_mha_partition[7].op == "MatMul"
):
fp8_mha_partitions.append(fp8_mha_partition)
print(f"Found {len(fp8_mha_partitions)} FP8 attentions")
# Insert Cast nodes for BMM1 and BMM2.
for fp8_mha_partition in fp8_mha_partitions:
bmm1_node = fp8_mha_partition[1]
insert_cast(graph, input_tensor=bmm1_node.inputs[0], attrs={"to": np.float32})
insert_cast(graph, input_tensor=bmm1_node.inputs[1], attrs={"to": np.float32})
insert_cast(graph, input_tensor=bmm1_node.outputs[0], attrs={"to": np.float16})
bmm2_node = fp8_mha_partition[7]
insert_cast(graph, input_tensor=bmm2_node.inputs[0], attrs={"to": np.float32})
insert_cast(graph, input_tensor=bmm2_node.inputs[1], attrs={"to": np.float32})
insert_cast(graph, input_tensor=bmm2_node.outputs[0], attrs={"to": np.float16})
def set_quant_precision(quant_config, precision: str = "Half"):
for key in quant_config["quant_cfg"]:
if "trt_high_precision_dtype" in quant_config["quant_cfg"][key]:
quant_config["quant_cfg"][key]["trt_high_precision_dtype"] = precision
def convert_fp16_io(graph):
"""
Convert graph I/O to FP16.
"""
for input_tensor in graph.inputs:
input_tensor.dtype = onnx.TensorProto.FLOAT16
for output_tensor in graph.outputs:
output_tensor.dtype = onnx.TensorProto.FLOAT16
def random_resize(cur_size: int):
"""
Randomly selects a new resolution for an image based on its current aspect ratio.
This function determines the current aspect ratio of an image, selects a new aspect ratio
from predefined choices depending on whether the current aspect ratio is square,
portrait, or landscape, and returns the corresponding resolution from a provided mapping.
Parameters:
cur_size (int): A tuple (width, height) representing the current resolution of the image.
resolution_to_aspects (dict[float, tuple[int, int]]): A mapping of aspect ratios (floats)
to their corresponding resolutions as tuples of (width, height).
Returns:
tuple[int, int]: A tuple (new_width, new_height) representing the newly selected resolution.
Raises:
KeyError: If the chosen aspect ratio is not present in the `resolution_to_aspects` dictionary.
Notes:
- For square images (aspect ratio = 1), the function selects from aspect ratios 1.25, 0.8, 1.5, and 0.667.
- For landscape images (aspect ratio > 1), the function selects from aspect ratios 1.778, 1.25, and 1.5.
- For portrait images (aspect ratio < 1), the function selects from aspect ratios 0.563, 0.8, and 0.667.
"""
resolution_to_aspects = {
1.0: (1024, 1024),
1.778: (768, 1344),
0.563: (1344, 768),
1.25: (896, 1152),
0.8: (1152, 896),
1.5: (832, 1216),
0.667: (1216, 832),
}
cur_aspect_ratio = round(cur_size[1] / cur_size[0], 3)
if cur_aspect_ratio == 1:
new_aspect_ratio = choice((1.25, 0.8, 1.5, 0.667))
new_res = resolution_to_aspects[new_aspect_ratio]
elif cur_aspect_ratio > 1:
new_aspect_ratio = choice((1.778, 1.25, 1.5))
new_res = resolution_to_aspects[new_aspect_ratio]
else:
# cur_aspect_ratio < 1
new_aspect_ratio = choice((0.563, 0.8, 0.667))
new_res = resolution_to_aspects[new_aspect_ratio]
return new_res
class PromptImageDataset(Dataset):
def __init__(
self,
root_dir,
):
"""
Args:
root_dir (str): Directory with all the images and the prompt file.
"""
self.root_dir = root_dir
self.possible_resolutions = {1024, 768, 1344, 896, 832, 1216}
self.global_idx_template = "{} | {} | {}"
self.prompts_by_size = defaultdict(list)
self.images_by_size = defaultdict(list)
self.images = []
self.prompts = []
self.images_size = []
# self.global_idx_2_group = dict()
# self.global_idx_to_group_idx = dict()
self.group_to_global_idx = {}
for idx, file in enumerate(os.listdir(os.path.join(self.root_dir, "prompts"))):
if not file.endswith(".txt"):
continue
file_name = os.path.splitext(file)[0]
image_path = os.path.join(
self.root_dir,
"inputs",
f"{file_name}.png",
)
with Image.open(image_path) as img, open(os.path.join(self.root_dir, "prompts", file), "r") as f:
prompt = "\n".join(f.readlines())
std_img_size = (
self.closest_value(img.size[0], self.possible_resolutions),
self.closest_value(img.size[1], self.possible_resolutions),
)
self.images_by_size[std_img_size].append(image_path)
self.prompts_by_size[std_img_size].append(prompt)
self.images.append(image_path)
self.prompts.append(prompt)
self.images_size.append(std_img_size)
# create a unique key that map group and index inside the group to a global index
in_group_idx = len(self.images_by_size[std_img_size]) - 1
group_idx_key = self.global_idx_template.format(std_img_size[0], std_img_size[1], in_group_idx)
self.group_to_global_idx[group_idx_key] = len(self.images) - 1
assert len(self.images) == len(self.prompts)
assert len(self.images) == len(self.group_to_global_idx)
@staticmethod
def closest_value(target: int, candidates: Set[int]):
"""
Find the closest value to the target from a set of candidate values.
Args:
target (int): The integer to compare against.
candidates (set): A set of integers as candidates.
Returns:
int: The closest value from the candidates.
"""
if not candidates:
raise ValueError("The candidates set cannot be empty.")
# Use the min function with a key that computes the absolute difference
return min(candidates, key=lambda x: abs(x - target))
def __len__(self):
return len(self.images)
def __getitem__(self, idx):
"""
Returns:
image (Tensor): Transformed image.
prompt (str): Corresponding text prompt.
"""
if torch.is_tensor(idx):
idx = idx.tolist()
prompt = self.prompts[idx]
image = self.images[idx]
image_size = self.images_size[idx]
return image, prompt, image_size
class SameSizeSampler(Sampler):
def __init__(self, dataset: PromptImageDataset, batch_size: int):
"""
Custom sampler that creates batches of images with the same size
Args:
dataset (SameSizeImageDataset): Dataset to sample from
batch_size (int): Number of images per batch
"""
super().__init__(dataset)
self.dataset = dataset
self.batch_size = batch_size
# Prepare size groups with indices
self.size_groups = {}
for size, image_paths in self.dataset.images_by_size.items():
# Create a list of indices for this size group
self.size_groups[size] = list(range(len(image_paths)))
def __iter__(self):
"""
Iteration method that yields indices for batches of same-size images
"""
# Create a copy of size groups to shuffle
size_groups_copy = {std_img_size: indices.copy() for std_img_size, indices in self.size_groups.items()}
# Shuffle each size group
for std_img_size, indices in size_groups_copy.items():
shuffle(indices)
# Iterate through size groups
for std_img_size, indices in size_groups_copy.items():
# Batch indices of the same size
for i in range(0, len(indices), self.batch_size):
# Yield batch indices for this size
batch_group_idxs = indices[i : min(i + self.batch_size, len(indices))]
for in_group_idx in batch_group_idxs:
group_idx_key = self.dataset.global_idx_template.format(
std_img_size[0], std_img_size[1], in_group_idx
)
batch_global_idx = self.dataset.group_to_global_idx[group_idx_key]
# batch_global_idxs.append(batch_global_idx)
yield batch_global_idx
def __len__(self):
"""
Total number of batches
"""
return len(self.dataset.images) // self.batch_size
def custom_collate(data):
"""
Custom collate function to handle batches of same-size images
Args:
dataset (SameSizeImageDataset): Dataset instance
batch (list): List of global indices
Returns:
tuple: Batched images and their size
"""
# Group images by their size
images, prompts, image_sizes = tuple(map(list, zip(*data)))
assert len(images) > 0
new_img_size = random_resize(image_sizes[0])
batch_images = []
for image in images:
with Image.open(image) as image:
image = image.convert("RGB").resize(size=new_img_size, resample=Image.LANCZOS)
image = np.array(image)
image = np.transpose(image, axes=(-1, 0, 1))
image = torch.from_numpy(image).float() / 127.5 - 1.0
batch_images.append(image)
batch_images = torch.stack(batch_images, dim=0)
return batch_images, prompts
def infinite_dataloader(dataloader):
while True:
for batch in dataloader:
yield batch