diff --git a/installer.py b/installer.py index d59908bf0..c0fd2faf9 100644 --- a/installer.py +++ b/installer.py @@ -603,7 +603,7 @@ def check_diffusers(): if args.skip_git: install('diffusers') return - sha = '9b721db205729d5a6e97a72312c3a0f4534064f1' # diffusers commit hash + sha = '764b62473ac10afd9c52e6b3f3f528f719bc7a34' # diffusers commit hash pkg = pkg_resources.working_set.by_key.get('diffusers', None) minor = int(pkg.version.split('.')[1] if pkg is not None else -1) cur = opts.get('diffusers_version', '') if minor > -1 else '' diff --git a/modules/model_quant.py b/modules/model_quant.py index 3275a6e91..45f940c59 100644 --- a/modules/model_quant.py +++ b/modules/model_quant.py @@ -12,6 +12,7 @@ from installer import installed, install, log, setup_logging ao = None bnb = None optimum_quanto = None +trt = None quant_last_model_name = None quant_last_model_device = None debug = os.environ.get('SD_QUANT_DEBUG', None) is not None @@ -103,6 +104,33 @@ def create_quanto_config(kwargs = None, allow: bool = True, module: str = 'Model return kwargs +def create_trt_config(kwargs = None, allow: bool = True, module: str = 'Model', modules_to_not_convert: list = []): + from modules import shared + if allow and (module == 'any' or module in shared.opts.trt_quantization): + load_trt() + if trt is None: + return kwargs + trt_config_data = { + "fp8": {"quant_type": "FP8", "quant_method": "modelopt", "modules_to_not_convert": []}, + "int8": {"quant_type": "INT8", "quant_method": "modelopt", "modules_to_not_convert": []}, + "int4": {"quant_type": "INT4", "quant_method": "modelopt", "block_quantize": 128, "channel_quantize": -1, "modules_to_not_convert": ["conv", "patch_embed"]}, + "nf4": {"quant_type": "NF4", "quant_method": "modelopt", "block_quantize": 128, "channel_quantize": -1, "scale_block_quantize": 8, "scale_channel_quantize": -1, "modules_to_not_convert": ["conv"]}, + "nvfp4": {"quant_type": "NVFP4", "quant_method": "modelopt", "block_quantize": 128, "channel_quantize": -1, "modules_to_not_convert": ["conv"]}, + } + trt_quant_config = trt_config_data[shared.opts.trt_quantization_type].copy() + for m in modules_to_not_convert: + if m not in trt_quant_config['modules_to_not_convert']: + trt_quant_config['modules_to_not_convert'].append(m) + trt_config = diffusers.quantizers.quantization_config.NVIDIAModelOptConfig(**trt_quant_config) + log.debug(f'Quantization: module={module} type=tensorrt dtype={shared.opts.trt_quantization_type}') + if kwargs is None: + return trt_config + else: + kwargs['quantization_config'] = trt_config + return kwargs + return kwargs + + def get_sdnq_devices(): from modules import devices, shared if shared.opts.device_map == "gpu": @@ -122,6 +150,7 @@ def get_sdnq_devices(): return_device = None return quantization_device, return_device + def create_sdnq_config(kwargs = None, allow: bool = True, module: str = 'Model', weights_dtype: str = None, modules_to_not_convert: list = [], modules_dtype_dict: dict = {}): from modules import shared if allow and (shared.opts.sdnq_quantize_mode in {'pre', 'auto'}) and (module == 'any' or module in shared.opts.sdnq_quantize_weights): @@ -224,6 +253,11 @@ def create_config(kwargs = None, allow: bool = True, module: str = 'Model', modu if debug: log.trace(f'Quantization: type=torchao config={kwargs.get("quantization_config", None)}') return kwargs + kwargs = create_trt_config(kwargs, allow=allow, module=module, modules_to_not_convert=modules_to_not_convert) + if kwargs is not None and 'quantization_config' in kwargs: + if debug: + log.trace(f'Quantization: type=tensorrt config={kwargs.get("quantization_config", None)}') + return kwargs return kwargs @@ -309,6 +343,30 @@ def load_quanto(msg='', silent=False): return None +def load_trt(msg='', silent=False): + global trt # pylint: disable=global-statement + if trt is not None: + return trt + try: + install('nvidia-modelopt') + import pydantic + if pydantic.__version__.startswith('1'): + log.error('Quantization: type=tensorrt pydantic==2 required') + return None + import modelopt + trt = modelopt + fn = f'{sys._getframe(3).f_code.co_name}:{sys._getframe(2).f_code.co_name}:{sys._getframe(1).f_code.co_name}' # pylint: disable=protected-access + log.debug(f'Quantization: type=tensorrt version={trt.__version__} fn={fn}') # pylint: disable=protected-access + return trt + except Exception as e: + if len(msg) > 0: + log.error(f"{msg} failed to import tensorrt: {e}") + trt = None + if not silent: + raise + return None + + def upcast_non_layerwise_modules(model, dtype): # pylint: disable=unused-argument from diffusers.hooks.layerwise_casting import _GO_LC_SUPPORTED_PYTORCH_LAYERS model_children = list(model.children()) diff --git a/modules/shared.py b/modules/shared.py index 0e67b5fd1..f5774461a 100644 --- a/modules/shared.py +++ b/modules/shared.py @@ -193,6 +193,11 @@ options_templates.update(options_section(("quantization", "Model Quantization"), "sdnq_dequantize_fp32": OptionInfo(False, "Dequantize using full precision", gr.Checkbox), "sdnq_quantize_shuffle_weights": OptionInfo(False, "Shuffle weights in post mode", gr.Checkbox), + "nunchaku_sep": OptionInfo("

Nunchaku Engine

", "", gr.HTML), + "nunchaku_quantization": OptionInfo([], "SVDQuant enabled", gr.CheckboxGroup, {"choices": ["Model", "TE"]}), + "nunchaku_attention": OptionInfo(False, "Nunchaku attention", gr.Checkbox), + "nunchaku_offload": OptionInfo(False, "Nunchaku offloading", gr.Checkbox), + "bnb_quantization_sep": OptionInfo("

BitsAndBytes

", "", gr.HTML), "bnb_quantization": OptionInfo([], "Quantization enabled", gr.CheckboxGroup, {"choices": ["Model", "TE", "LLM", "VAE"]}), "bnb_quantization_type": OptionInfo("nf4", "Quantization type", gr.Dropdown, {"choices": ["nf4", "fp8", "fp4"]}), @@ -218,10 +223,9 @@ options_templates.update(options_section(("quantization", "Model Quantization"), "layerwise_quantization_storage": OptionInfo("float8_e4m3fn", "Layerwise casting storage", gr.Dropdown, {"choices": ["float8_e4m3fn", "float8_e5m2"]}), "layerwise_quantization_nonblocking": OptionInfo(False, "Layerwise non-blocking operations", gr.Checkbox), - "nunchaku_sep": OptionInfo("

Nunchaku Engine

", "", gr.HTML), - "nunchaku_quantization": OptionInfo([], "SVDQuant enabled", gr.CheckboxGroup, {"choices": ["Model", "TE"]}), - "nunchaku_attention": OptionInfo(False, "Nunchaku attention", gr.Checkbox), - "nunchaku_offload": OptionInfo(False, "Nunchaku offloading", gr.Checkbox), + "trt_quantization_sep": OptionInfo("

TensorRT

", "", gr.HTML), + "trt_quantization": OptionInfo([], "Quantization enabled", gr.CheckboxGroup, {"choices": ["Model", "TE", "LLM", "VAE"]}), + "trt_quantization_type": OptionInfo("nvfp4", "Quantization type", gr.Dropdown, {"choices": ["int8", "int4", "fp8", "nf4", "nvfp4"]}), "nncf_compress_sep": OptionInfo("

NNCF: Neural Network Compression Framework

", "", gr.HTML, {"visible": cmd_opts.use_openvino}), "nncf_compress_weights": OptionInfo([], "Quantization enabled", gr.CheckboxGroup, {"choices": ["Model", "TE", "VAE"], "visible": cmd_opts.use_openvino}),