wehub-resource-sync
466 lines
20 KiB
Python
466 lines
20 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.
|
|
#
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
from typing import Any, Dict, Tuple
|
|
|
|
import torch
|
|
|
|
# Define valid optimization levels for TensorRT engine build
|
|
VALID_OPTIMIZATION_LEVELS = list(range(6))
|
|
|
|
|
|
def parse_key_value_pairs(string: str) -> Dict[str, str]:
|
|
"""Parse a string of comma-separated key-value pairs into a dictionary.
|
|
|
|
Args:
|
|
string (str): A string of comma-separated key-value pairs.
|
|
|
|
Returns:
|
|
Dict[str, str]: Parsed dictionary of key-value pairs.
|
|
|
|
Example:
|
|
>>> parse_key_value_pairs("key1:value1,key2:value2")
|
|
{"key1": "value1", "key2": "value2"}
|
|
"""
|
|
parsed = {}
|
|
|
|
for key_value_pair in string.split(","):
|
|
if not key_value_pair:
|
|
continue
|
|
|
|
key_value_pair = key_value_pair.split(":")
|
|
if len(key_value_pair) != 2:
|
|
raise argparse.ArgumentTypeError(f"Invalid key-value pair: {key_value_pair}. Must have length 2.")
|
|
key, value = key_value_pair
|
|
parsed[key] = value
|
|
|
|
return parsed
|
|
|
|
|
|
def add_arguments(parser):
|
|
# Stable Diffusion configuration
|
|
parser.add_argument(
|
|
"--version",
|
|
type=str,
|
|
default="1.4",
|
|
choices=(
|
|
"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-schnell",
|
|
"flux.1-dev-canny",
|
|
"flux.1-dev-depth",
|
|
"flux.1-kontext-dev",
|
|
"cosmos-predict2-2b-text2image",
|
|
"cosmos-predict2-14b-text2image",
|
|
"cosmos-predict2-2b-video2world",
|
|
"cosmos-predict2-14b-video2world",
|
|
),
|
|
help="Version of Stable Diffusion",
|
|
)
|
|
parser.add_argument("prompt", nargs="*", help="Text prompt(s) to guide image generation")
|
|
parser.add_argument(
|
|
"--negative-prompt", nargs="*", default=[""], help="The negative prompt(s) to guide the image generation."
|
|
)
|
|
parser.add_argument("--batch-size", type=int, default=1, choices=[1, 2, 4], help="Batch size (repeat prompt)")
|
|
parser.add_argument(
|
|
"--batch-count", type=int, default=1, help="Number of images to generate in sequence, one at a time."
|
|
)
|
|
parser.add_argument("--height", type=int, default=512, help="Height of image to generate (must be multiple of 8)")
|
|
parser.add_argument("--width", type=int, default=512, help="Height of image to generate (must be multiple of 8)")
|
|
parser.add_argument("--denoising-steps", type=int, default=30, help="Number of denoising steps")
|
|
parser.add_argument(
|
|
"--scheduler",
|
|
type=str,
|
|
default=None,
|
|
choices=("DDIM", "DDPM", "EulerA", "Euler", "LCM", "LMSD", "PNDM", "UniPC", "DDPMWuerstchen", "FlowMatchEuler"),
|
|
help="Scheduler for diffusion process",
|
|
)
|
|
parser.add_argument(
|
|
"--guidance-scale",
|
|
type=float,
|
|
default=7.5,
|
|
help="Value of classifier-free guidance scale (must be greater than 1)",
|
|
)
|
|
parser.add_argument(
|
|
"--lora-scale",
|
|
type=float,
|
|
default=1.0,
|
|
help="Controls how much to influence the outputs with the LoRA parameters. (must between 0 and 1)",
|
|
)
|
|
parser.add_argument(
|
|
"--lora-weight",
|
|
type=float,
|
|
nargs="+",
|
|
default=None,
|
|
help="The LoRA adapter(s) weights to use with the UNet. (must between 0 and 1)",
|
|
)
|
|
parser.add_argument(
|
|
"--lora-path",
|
|
type=str,
|
|
nargs="+",
|
|
default=None,
|
|
help="Path to LoRA adaptor. Ex: 'latent-consistency/lcm-lora-sdv1-5'",
|
|
)
|
|
parser.add_argument("--bf16", action="store_true", help="Run pipeline in BFloat16 precision")
|
|
|
|
# ONNX export
|
|
parser.add_argument(
|
|
"--onnx-opset",
|
|
type=int,
|
|
default=19,
|
|
choices=range(7, 24),
|
|
help="Select ONNX opset version to target for exported models",
|
|
)
|
|
parser.add_argument("--onnx-dir", default="onnx", help="Output directory for ONNX export")
|
|
parser.add_argument(
|
|
"--custom-onnx-paths",
|
|
type=parse_key_value_pairs,
|
|
help=(
|
|
"[FLUX, Stable Diffusion 3.5-large, Cosmos only] Custom override paths to pre-exported ONNX model files. These ONNX models are directly used to "
|
|
"build TRT engines without further optimization on the ONNX graphs. Paths should be a comma-separated list "
|
|
"of <model_name>:<path> pairs. For example: "
|
|
"--custom-onnx-paths=transformer:/path/to/transformer.onnx,vae:/path/to/vae.onnx. Call "
|
|
"<PipelineClass>.get_model_names(...) for the list of supported model names."
|
|
),
|
|
)
|
|
parser.add_argument(
|
|
"--onnx-export-only",
|
|
action="store_true",
|
|
help="If set, only performs the export of models to ONNX, skipping engine build and inference.",
|
|
)
|
|
parser.add_argument(
|
|
"--download-onnx-models",
|
|
action="store_true",
|
|
help=("[FLUX and Stable Diffusion 3.5-large only] Download pre-exported ONNX models"),
|
|
)
|
|
|
|
# Framework model ckpt
|
|
parser.add_argument("--framework-model-dir", default="pytorch_model", help="Directory for HF saved models")
|
|
|
|
# TensorRT engine build
|
|
parser.add_argument("--engine-dir", default="engine", help="Output directory for TensorRT engines")
|
|
parser.add_argument(
|
|
"--custom-engine-paths",
|
|
type=parse_key_value_pairs,
|
|
help=(
|
|
"[FLUX only] Custom override paths to pre-built engine files. Paths should be a comma-separated list of "
|
|
"<model_name>:<path> pairs. For example: "
|
|
"--custom-onnx-paths=transformer:/path/to/transformer.plan,vae:/path/to/vae.plan. Call "
|
|
"<PipelineClass>.get_model_names(...) for the list of supported model names."
|
|
),
|
|
)
|
|
|
|
parser.add_argument(
|
|
"--optimization-level",
|
|
type=int,
|
|
default=None,
|
|
help=f"Set the builder optimization level to build the engine with. A higher level allows TensorRT to spend more building time for more optimization options. Must be one of {VALID_OPTIMIZATION_LEVELS}.",
|
|
)
|
|
parser.add_argument(
|
|
"--build-static-batch", action="store_true", help="Build TensorRT engines with fixed batch size."
|
|
)
|
|
parser.add_argument(
|
|
"--build-dynamic-shape", action="store_true", help="Build TensorRT engines with dynamic image shapes."
|
|
)
|
|
parser.add_argument(
|
|
"--build-enable-refit", action="store_true", help="Enable Refit option in TensorRT engines during build."
|
|
)
|
|
parser.add_argument(
|
|
"--build-all-tactics", action="store_true", help="Build TensorRT engines using all tactic sources."
|
|
)
|
|
parser.add_argument(
|
|
"--timing-cache", default=None, type=str, help="Path to the precached timing measurements to accelerate build."
|
|
)
|
|
parser.add_argument("--ws", action="store_true", help="Build TensorRT engines with weight streaming enabled.")
|
|
|
|
# Quantization configuration.
|
|
parser.add_argument("--int8", action="store_true", help="Apply int8 quantization.")
|
|
parser.add_argument("--fp8", action="store_true", help="Apply fp8 quantization.")
|
|
parser.add_argument("--fp4", action="store_true", help="Apply fp4 quantization.")
|
|
parser.add_argument(
|
|
"--quantization-level",
|
|
type=float,
|
|
default=0.0,
|
|
choices=[0.0, 1.0, 2.0, 2.5, 3.0, 4.0],
|
|
help="int8/fp8 quantization level, 1: CNN, 2: CNN + FFN, 2.5: CNN + FFN + QKV, 3: CNN + Almost all Linear (Including FFN, QKV, Proj and others), 4: CNN + Almost all Linear + fMHA, 0: Default to 2.5 for int8 and 4.0 for fp8.",
|
|
)
|
|
parser.add_argument(
|
|
"--quantization-percentile",
|
|
type=float,
|
|
default=1.0,
|
|
help="Control quantization scaling factors (amax) collecting range, where the minimum amax in range(n_steps * percentile) will be collected. Recommendation: 1.0.",
|
|
)
|
|
parser.add_argument(
|
|
"--quantization-alpha",
|
|
type=float,
|
|
default=0.8,
|
|
help="The alpha parameter for SmoothQuant quantization used for linear layers. Recommendation: 0.8 for SDXL.",
|
|
)
|
|
parser.add_argument(
|
|
"--calibration-size",
|
|
type=int,
|
|
default=32,
|
|
help="The number of steps to use for calibrating the model for quantization. Recommendation: 32, 64, 128 for SDXL",
|
|
)
|
|
|
|
# Inference
|
|
parser.add_argument(
|
|
"--num-warmup-runs", type=int, default=5, help="Number of warmup runs before benchmarking performance"
|
|
)
|
|
parser.add_argument("--use-cuda-graph", action="store_true", help="Enable cuda graph")
|
|
parser.add_argument("--nvtx-profile", action="store_true", help="Enable NVTX markers for performance profiling")
|
|
parser.add_argument(
|
|
"--torch-inference",
|
|
default="",
|
|
help="Run inference with PyTorch (using specified compilation mode) instead of TensorRT.",
|
|
)
|
|
parser.add_argument(
|
|
"--torch-fallback",
|
|
default=None,
|
|
type=str,
|
|
help="[FLUX, SD3.5, and Wan] Comma separated list of models to be inferenced using PyTorch instead of TRT. For example --torch-fallback text_encoder,transformer,transformer_2. If --torch-inference set, this parameter will be ignored.",
|
|
)
|
|
parser.add_argument(
|
|
"--low-vram",
|
|
action="store_true",
|
|
help="[FLUX, SD3.5, and Wan] Optimize for low VRAM usage, possibly at the expense of inference performance. Disabled by default.",
|
|
)
|
|
parser.add_argument("--seed", type=int, default=None, help="Seed for random generator to get consistent results")
|
|
parser.add_argument("--output-dir", default="output", help="Output directory for logs and image artifacts")
|
|
parser.add_argument("--hf-token", type=str, help="HuggingFace API access token for downloading model checkpoints")
|
|
parser.add_argument("-v", "--verbose", action="store_true", help="Show verbose output")
|
|
return parser
|
|
|
|
|
|
def process_pipeline_args(args: argparse.Namespace) -> Tuple[Dict[str, Any], Dict[str, Any], Tuple]:
|
|
"""Validate parsed arguments and process argument values.
|
|
|
|
Some argument values are resolved or overwritten during processing.
|
|
|
|
Args:
|
|
args (argparse.Namespace): Parsed argument. This is modified in-place.
|
|
|
|
Returns:
|
|
Dict[str, Any]: Keyword arguments for initializing a pipeline. This is only used in legacy pipelines that do not
|
|
have factory methods `FromArgs` that construct the pipeline directly from the parsed argument.
|
|
Dict[str, Any]: Keyword arguments for calling the `.load_engine` method of the pipeline.
|
|
Tuple: Arguments for calling the `.run` method of the pipeline.
|
|
"""
|
|
|
|
# GPU device info
|
|
device_info = torch.cuda.get_device_properties(0)
|
|
sm_version = device_info.major * 10 + device_info.minor
|
|
|
|
is_flux = args.version.startswith("flux")
|
|
is_sd35 = args.version.startswith("3.5")
|
|
is_wan = args.version.startswith("wan")
|
|
is_cosmos = args.version.startswith("cosmos")
|
|
|
|
if args.height % 8 != 0 or args.width % 8 != 0:
|
|
raise ValueError(
|
|
f"Image height and width have to be divisible by 8 but specified as: {args.image_height} and {args.width}."
|
|
)
|
|
|
|
# Handle batch size
|
|
max_batch_size = 4
|
|
if args.batch_size > max_batch_size:
|
|
raise ValueError(f"Batch size {args.batch_size} is larger than allowed {max_batch_size}.")
|
|
|
|
if args.use_cuda_graph and (not args.build_static_batch or args.build_dynamic_shape):
|
|
raise ValueError(
|
|
"Using CUDA graph requires static dimensions. Enable `--build-static-batch` and do not specify `--build-dynamic-shape`"
|
|
)
|
|
|
|
# TensorRT builder optimization level
|
|
if args.optimization_level is None:
|
|
# optimization level set to 3 for all Flux pipelines to reduce GPU memory usage
|
|
if args.int8 or args.fp8 and not is_flux:
|
|
args.optimization_level = 4
|
|
else:
|
|
args.optimization_level = 3
|
|
|
|
if args.optimization_level not in VALID_OPTIMIZATION_LEVELS:
|
|
raise ValueError(
|
|
f"Optimization level {args.optimization_level} not valid. Valid values are: {VALID_OPTIMIZATION_LEVELS}"
|
|
)
|
|
|
|
# Quantized pipeline
|
|
# int8 support
|
|
if args.int8 and not any(args.version.startswith(prefix) for prefix in ("xl", "1.4")):
|
|
raise ValueError("int8 quantization is only supported for SDXL and SD1.4 pipelines.")
|
|
|
|
# fp8 support validation
|
|
if args.fp8:
|
|
# Check version compatibility
|
|
supported_versions = ("xl", "1.4", "3.5-large")
|
|
if not (any(args.version.startswith(prefix) for prefix in supported_versions) or is_flux):
|
|
raise ValueError(
|
|
"fp8 quantization is only supported for SDXL, SD1.4, SD3.5-large and FLUX pipelines."
|
|
)
|
|
|
|
# Check controlnet compatibility
|
|
if getattr(args, "controlnet_type", None) is not None:
|
|
if args.version not in ("xl-1.0", "3.5-large"):
|
|
raise ValueError("fp8 controlnet quantization is only supported for SDXL and SD3.5-large.")
|
|
if args.version == "3.5-large" and args.controlnet_type == "blur":
|
|
raise ValueError("Blur controlnet type is not supported for SD3.5.")
|
|
# Check for conflicting quantization
|
|
if args.int8:
|
|
raise ValueError("Cannot apply both int8 and fp8 quantization, please choose only one.")
|
|
|
|
# Check GPU compute capability
|
|
if sm_version < 89:
|
|
raise ValueError(
|
|
f"Cannot apply FP8 quantization for GPU with compute capability {sm_version / 10.0}. A minimum compute capability of 8.9 is required."
|
|
)
|
|
|
|
# Check SD3.5-large specific requirement
|
|
if args.version == "3.5-large" and not args.download_onnx_models:
|
|
raise ValueError(
|
|
"Native FP8 quantization is not supported for SD3.5-large. Please pass --download-onnx-models."
|
|
)
|
|
|
|
# TensorRT ModelOpt quantization level
|
|
if args.quantization_level == 0.0:
|
|
def override_quant_level(level: float, dtype_str: str):
|
|
args.quantization_level = level
|
|
print(f"[W] The default quantization level has been set to {level} for {dtype_str}.")
|
|
|
|
if args.fp8:
|
|
# L4 fp8 fMHA on Hopper not yet enabled.
|
|
if sm_version == 90 and is_flux:
|
|
override_quant_level(3.0, "FP8")
|
|
else:
|
|
override_quant_level(3.0 if args.version == "1.4" else 4.0, "FP8")
|
|
|
|
elif args.int8:
|
|
override_quant_level(3.0, "INT8")
|
|
|
|
if args.version.startswith("flux") and args.quantization_level == 3.0 and args.download_onnx_models:
|
|
raise ValueError(
|
|
"Transformer ONNX model for Quantization level 3 is not available for download. Please export the quantized Transformer model natively with the removal of --download-onnx-models."
|
|
)
|
|
if args.fp4:
|
|
# FP4 precision is only supported for the Flux pipeline
|
|
assert is_flux, "FP4 precision is only supported for the Flux pipeline"
|
|
|
|
# Handle LoRA
|
|
# FLUX canny and depth official LoRAs are not supported because they modify the transformer architecture, conflicting with refit
|
|
if args.lora_path and not any(args.version.startswith(prefix) for prefix in ("xl", "flux.1-dev", "flux.1-schnell")):
|
|
raise ValueError("LoRA adapter support is only supported for SDXL, FLUX.1-dev and FLUX.1-schnell pipelines")
|
|
|
|
if args.lora_weight:
|
|
for weight in (weight for weight in args.lora_weight if not 0 <= weight <= 1):
|
|
raise ValueError(f"LoRA adapter weights must be between 0 and 1, provided {weight}")
|
|
|
|
if not 0 <= args.lora_scale <= 1:
|
|
raise ValueError(f"LoRA scale value must be between 0 and 1, provided {args.lora_scale}")
|
|
|
|
# Force lora merge when fp8 or int8 is used with LoRA
|
|
if args.build_enable_refit and args.lora_path and (args.int8 or args.fp8):
|
|
raise ValueError(
|
|
"Engine refit should not be enabled for quantized models with LoRA. ModelOpt recommends fusing the LoRA to the model before quantization. \
|
|
See https://github.com/NVIDIA/TensorRT-Model-Optimizer/tree/main/examples/diffusers/quantization#lora"
|
|
)
|
|
|
|
# Torch-fallback and Torch-inference
|
|
if args.torch_fallback and not args.torch_inference:
|
|
assert (
|
|
is_flux or is_sd35 or is_wan or is_cosmos
|
|
), "PyTorch Fallback is only supported for Flux, Stable Diffusion 3.5, Wan and Cosmos pipelines."
|
|
args.torch_fallback = args.torch_fallback.split(",")
|
|
|
|
if args.torch_fallback and args.torch_inference:
|
|
print(
|
|
"[W] All models will run in PyTorch when --torch-inference is set. Parameter --torch-fallback will be ignored."
|
|
)
|
|
args.torch_fallback = None
|
|
|
|
# low-vram
|
|
if args.low_vram:
|
|
assert (
|
|
is_flux or is_sd35 or is_wan or is_cosmos
|
|
), "low-vram mode is only supported for Flux, Stable Diffusion 3.5, Wan and Cosmos pipelines."
|
|
|
|
# Disable SDXL LCM pipeline
|
|
if args.version == "xl-1.0" and args.scheduler == "LCM":
|
|
raise ValueError("SDXL pipeline does not support the LCM scheduler currently. Please use a different scheduler.")
|
|
|
|
# Pack arguments
|
|
kwargs_init_pipeline = {
|
|
"version": args.version,
|
|
"max_batch_size": max_batch_size,
|
|
"denoising_steps": args.denoising_steps,
|
|
"scheduler": args.scheduler,
|
|
"guidance_scale": args.guidance_scale,
|
|
"output_dir": args.output_dir,
|
|
"hf_token": args.hf_token,
|
|
"verbose": args.verbose,
|
|
"nvtx_profile": args.nvtx_profile,
|
|
"use_cuda_graph": args.use_cuda_graph,
|
|
"lora_scale": args.lora_scale,
|
|
"lora_weight": args.lora_weight,
|
|
"lora_path": args.lora_path,
|
|
"framework_model_dir": args.framework_model_dir,
|
|
"torch_inference": args.torch_inference,
|
|
}
|
|
|
|
kwargs_load_engine = {
|
|
"onnx_opset": args.onnx_opset,
|
|
"opt_batch_size": args.batch_size,
|
|
"opt_image_height": args.height,
|
|
"opt_image_width": args.width,
|
|
"optimization_level": args.optimization_level,
|
|
"static_batch": args.build_static_batch,
|
|
"static_shape": not args.build_dynamic_shape,
|
|
"enable_all_tactics": args.build_all_tactics,
|
|
"enable_refit": args.build_enable_refit,
|
|
"timing_cache": args.timing_cache,
|
|
"int8": args.int8,
|
|
"fp8": args.fp8,
|
|
"fp4": args.fp4,
|
|
"quantization_level": args.quantization_level,
|
|
"quantization_percentile": args.quantization_percentile,
|
|
"quantization_alpha": args.quantization_alpha,
|
|
"calibration_size": args.calibration_size,
|
|
"onnx_export_only": args.onnx_export_only,
|
|
"download_onnx_models": args.download_onnx_models,
|
|
}
|
|
|
|
args_run_demo = (
|
|
args.prompt,
|
|
args.negative_prompt,
|
|
args.height,
|
|
args.width,
|
|
args.batch_size,
|
|
args.batch_count,
|
|
args.num_warmup_runs,
|
|
args.use_cuda_graph,
|
|
)
|
|
|
|
return kwargs_init_pipeline, kwargs_load_engine, args_run_demo
|