diff --git a/CHANGELOG.md b/CHANGELOG.md index bfd14cacd..a1b41e7a6 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -101,6 +101,13 @@ ### Updates +- **Quantization** + - Add `TorchAO` *pre* (during load) and *post* (during execution) quantization + **torchao** supports 4 different int-based and 3 float-based quantization schemes + This is in addition to existing support for: + - `BitsAndBytes` with 3 float-based quantization schemes + - `Optimium.Quanto` with 3 int-based and 2 float-based quantizations schemes + - `GGUF` with pre-quantized weights - **IPEX**: update to IPEX 2.5.10+xpu - **OpenVINO**: update to 2024.5.0 - **Sampler** improvements diff --git a/TODO.md b/TODO.md index 996da5ad9..fafbea8c3 100644 --- a/TODO.md +++ b/TODO.md @@ -10,21 +10,14 @@ Main ToDo list can be found at [GitHub projects](https://github.com/users/vladma ## Future Candidates -- SD35 IPAdapter: -- SD35 LoRA: -- Flux IPAdapter: +- SD35 IPAdapter: +- Flux IPAdapter: - Flux NF4: -- SANA: - LTX-Video: -- TorchAO: +- GGUF: ## Other - IPAdapter negative: - Control API enhance scripts compatibility - PixelSmith: - -## Workaround in place - -- GGUF -- FlowMatch diff --git a/installer.py b/installer.py index 18f2b9399..2bb3df492 100644 --- a/installer.py +++ b/installer.py @@ -459,7 +459,7 @@ def check_python(supported_minors=[9, 10, 11, 12], reason=None): def check_diffusers(): if args.skip_all or args.skip_requirements: return - sha = '5ed761a6f2a6dad56031f4e3e32223bfbe2dda01' # diffusers commit hash + sha = '1524781b88ac1a082e755a030ba9d73cd6948e84' # 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 0) cur = opts.get('diffusers_version', '') if minor > 0 else '' @@ -483,6 +483,7 @@ def check_onnx(): def check_torchao(): + """ if args.skip_all or args.skip_requirements: return if installed('torchao', quiet=True): @@ -492,6 +493,8 @@ def check_torchao(): pip('uninstall --yes torchao', ignore=True, quiet=True, uv=False) for m in [m for m in sys.modules if m.startswith('torchao')]: del sys.modules[m] + """ + return def install_cuda(): diff --git a/modules/model_flux.py b/modules/model_flux.py index 8d6a02ef6..362f96d9b 100644 --- a/modules/model_flux.py +++ b/modules/model_flux.py @@ -197,6 +197,7 @@ def load_transformer(file_path): # triggered by opts.sd_unet change transformer = _transformer else: diffusers_load_config = model_quant.create_bnb_config(diffusers_load_config) + diffusers_load_config = model_quant.create_ao_config(diffusers_load_config) transformer = diffusers.FluxTransformer2DModel.from_single_file(file_path, **diffusers_load_config) if transformer is None: shared.log.error('Failed to load UNet model') @@ -322,8 +323,9 @@ def load_flux(checkpoint_info, diffusers_load_config): # triggered by opts.sd_ch shared.log.warning(f'Load model: type=FLUX component={c} dtype={kwargs[c].dtype} cast dtype={devices.dtype} recast') kwargs[c] = kwargs[c].to(dtype=devices.dtype) - allow_bnb = 'gguf' not in (sd_unet.loaded_unet or '') - kwargs = model_quant.create_bnb_config(kwargs, allow_bnb) + allow_quant = 'gguf' not in (sd_unet.loaded_unet or '') + kwargs = model_quant.create_bnb_config(kwargs, allow_quant) + kwargs = model_quant.create_ao_config(kwargs, allow_quant) if checkpoint_info.path.endswith('.safetensors') and os.path.isfile(checkpoint_info.path): pipe = diffusers.FluxPipeline.from_single_file(checkpoint_info.path, cache_dir=shared.opts.diffusers_dir, **kwargs, **diffusers_load_config) else: diff --git a/modules/model_quant.py b/modules/model_quant.py index 5c0b40080..10e5eb99e 100644 --- a/modules/model_quant.py +++ b/modules/model_quant.py @@ -38,9 +38,9 @@ def create_ao_config(kwargs = None, allow_ao: bool = True): load_torchao() if ao is None: return kwargs - ao_config = {} - # ao_config = diffusers.TorchAoConfig("int8wo") # TODO torchao - shared.log.debug(f'Quantization: module=all type=bnb dtype={shared.opts.torchao_quantization_type}') + diffusers.utils.import_utils.is_torchao_available = lambda: True + ao_config = diffusers.TorchAoConfig(shared.opts.torchao_quantization_type) + shared.log.debug(f'Quantization: module=all type=torchao dtype={shared.opts.torchao_quantization_type}') if kwargs is None: return ao_config else: @@ -53,7 +53,7 @@ def load_torchao(msg='', silent=False): global ao # pylint: disable=global-statement if ao is not None: return ao - install('torchao', quiet=True) + install('torchao==0.7.0', quiet=True) try: import torchao ao = torchao diff --git a/modules/model_sana.py b/modules/model_sana.py index f9986d726..c25a7ffb9 100644 --- a/modules/model_sana.py +++ b/modules/model_sana.py @@ -38,6 +38,7 @@ def load_sana(checkpoint_info, kwargs={}): kwargs['variant'] = 'fp16' kwargs = model_quant.create_bnb_config(kwargs) + kwargs = model_quant.create_ao_config(kwargs) shared.log.debug(f'Load model: type=Sana repo="{repo_id}" args={kwargs}') t0 = time.time() pipe = diffusers.SanaPipeline.from_pretrained(repo_id, cache_dir = shared.opts.diffusers_dir, **kwargs) diff --git a/modules/model_tools.py b/modules/model_tools.py index 07cd61b6e..fdeda5c2a 100644 --- a/modules/model_tools.py +++ b/modules/model_tools.py @@ -69,11 +69,13 @@ def load_modules(repo_id: str, params: dict): subfolder = 'text_encoder_2' if cls == transformers.T5EncoderModel: # t5-xxl subfolder = 'text_encoder_3' - kwargs['quantization_config'] = model_quant.create_bnb_config() + kwargs = model_quant.create_bnb_config(kwargs) + kwargs = model_quant.create_ao_config(kwargs) kwargs['variant'] = 'fp16' if cls == diffusers.SD3Transformer2DModel: subfolder = 'transformer' - kwargs['quantization_config'] = model_quant.create_bnb_config() + kwargs = model_quant.create_bnb_config(kwargs) + kwargs = model_quant.create_ao_config(kwargs) if subfolder is None: continue shared.log.debug(f'Load: module={name} class={cls.__name__} repo={repo_id} location={subfolder}') diff --git a/modules/sd_models.py b/modules/sd_models.py index 661559ae9..52d0a1deb 100644 --- a/modules/sd_models.py +++ b/modules/sd_models.py @@ -282,7 +282,7 @@ def set_diffuser_options(sd_model, vae = None, op: str = 'model', offload=True): model.eval() return model sd_model = sd_models_compile.apply_compile_to_model(sd_model, eval_model, ["Model", "VAE", "Text Encoder"], op="eval") - if len(shared.opts.torchao_quantization) > 0 and shared.opts.torchao_quantization_mode != 'post': + if len(shared.opts.torchao_quantization) > 0 and shared.opts.torchao_quantization_mode == 'post': sd_model = sd_models_compile.torchao_quantization(sd_model) if shared.opts.opt_channelslast and hasattr(sd_model, 'unet'): diff --git a/modules/sd_models_compile.py b/modules/sd_models_compile.py index 38d3ef57f..bdd47e2e1 100644 --- a/modules/sd_models_compile.py +++ b/modules/sd_models_compile.py @@ -505,24 +505,13 @@ def compile_diffusers(sd_model): def torchao_quantization(sd_model): try: - install('torchao', quiet=True) + install('torchao==0.7.0', quiet=True) from torchao import quantization as q except Exception as e: shared.log.error(f"Quantization: type=TorchAO quantization not supported: {e}") return sd_model - if shared.opts.torchao_quantization_type == "int8+act": - fn = q.int8_dynamic_activation_int8_weight - elif shared.opts.torchao_quantization_type == "int8": - fn = q.int8_weight_only - elif shared.opts.torchao_quantization_type == "int4": - fn = q.int4_weight_only - elif shared.opts.torchao_quantization_type == "fp8+act": - fn = q.float8_dynamic_activation_float8_weight - elif shared.opts.torchao_quantization_type == "fp8": - fn = q.float8_weight_only - elif shared.opts.torchao_quantization_type == "fpx": - fn = q.fpx_weight_only - else: + fn = getattr(q, shared.opts.torchao_quantization_type, None) + if fn is None: shared.log.error(f"Quantization: type=TorchAO type={shared.opts.torchao_quantization_type} not supported") return sd_model shared.log.info(f"Quantization: type=TorchAO pipe={sd_model.__class__.__name__} quant={shared.opts.torchao_quantization_type} fn={fn} targets={shared.opts.torchao_quantization}") diff --git a/modules/shared.py b/modules/shared.py index 5e353b1b1..2a1e8c19d 100644 --- a/modules/shared.py +++ b/modules/shared.py @@ -583,23 +583,23 @@ options_templates.update(options_section(('backends', "Backend Settings"), { options_templates.update(options_section(('quantization', "Quantization Settings"), { "bnb_sep": OptionInfo("

BitsAndBytes

", "", gr.HTML), - "bnb_quantization": OptionInfo([], "Enabled", gr.CheckboxGroup, {"choices": ["Model", "VAE", "Text Encoder"], "visible": native}), - "bnb_quantization_type": OptionInfo("nf4", "Type", gr.Radio, {"choices": ['nf4', 'fp8', 'fp4'], "visible": native}), - "bnb_quantization_storage": OptionInfo("uint8", "Backend storage", gr.Radio, {"choices": ["float16", "float32", "int8", "uint8", "float64", "bfloat16"], "visible": native}), + "bnb_quantization": OptionInfo([], "Quantization enabled", gr.CheckboxGroup, {"choices": ["Model", "VAE", "Text Encoder"], "visible": native}), + "bnb_quantization_type": OptionInfo("nf4", "Quantization type", gr.Dropdown, {"choices": ['nf4', 'fp8', 'fp4'], "visible": native}), + "bnb_quantization_storage": OptionInfo("uint8", "Backend storage", gr.Dropdown, {"choices": ["float16", "float32", "int8", "uint8", "float64", "bfloat16"], "visible": native}), "optimum_quanto_sep": OptionInfo("

Optimum Quanto

", "", gr.HTML), - "optimum_quanto_weights": OptionInfo([], "Enabled", gr.CheckboxGroup, {"choices": ["Model", "VAE", "Text Encoder", "ControlNet"], "visible": native}), - "optimum_quanto_weights_type": OptionInfo("qint8", "Type", gr.Radio, {"choices": ['qint8', 'qfloat8_e4m3fn', 'qfloat8_e5m2', 'qint4', 'qint2'], "visible": native}), - "optimum_quanto_activations_type": OptionInfo("none", "Activations ", gr.Radio, {"choices": ['none', 'qint8', 'qfloat8_e4m3fn', 'qfloat8_e5m2'], "visible": native}), + "optimum_quanto_weights": OptionInfo([], "Quantization enabled", gr.CheckboxGroup, {"choices": ["Model", "VAE", "Text Encoder", "ControlNet"], "visible": native}), + "optimum_quanto_weights_type": OptionInfo("qint8", "Quantization weights type", gr.Dropdown, {"choices": ['qint8', 'qfloat8_e4m3fn', 'qfloat8_e5m2', 'qint4', 'qint2'], "visible": native}), + "optimum_quanto_activations_type": OptionInfo("none", "Quantization activations type ", gr.Dropdown, {"choices": ['none', 'qint8', 'qfloat8_e4m3fn', 'qfloat8_e5m2'], "visible": native}), "torchao_sep": OptionInfo("

TorchAO

", "", gr.HTML), - "torchao_quantization": OptionInfo([], "Enabled", gr.CheckboxGroup, {"choices": ["Model", "VAE", "Text Encoder"], "visible": native}), - "torchao_quantization_mode": OptionInfo("pre", "Mode", gr.Radio, {"choices": ['pre', 'post'], "visible": native}), - "torchao_quantization_type": OptionInfo("int8", "Type", gr.Radio, {"choices": ["int8+act", "int8", "int4", "fp8+act", "fp8", "fpx"], "visible": native}), + "torchao_quantization": OptionInfo([], "Quantization enabled", gr.CheckboxGroup, {"choices": ["Model", "VAE", "Text Encoder"], "visible": native}), + "torchao_quantization_mode": OptionInfo("pre", "Quantization mode", gr.Dropdown, {"choices": ['pre', 'post'], "visible": native}), + "torchao_quantization_type": OptionInfo("int8_weight_only", "Quantization type", gr.Dropdown, {"choices": ['int4_weight_only', 'int8_dynamic_activation_int4_weight', 'int8_weight_only', 'int8_dynamic_activation_int8_weight', 'float8_weight_only', 'float8_dynamic_activation_float8_weight', 'float8_static_activation_float8_weight'], "visible": native}), "nncf_sep": OptionInfo("

NNCF

", "", gr.HTML), - "nncf_compress_weights": OptionInfo([], "Enabled", gr.CheckboxGroup, {"choices": ["Model", "VAE", "Text Encoder", "ControlNet"], "visible": native}), - "nncf_compress_weights_mode": OptionInfo("INT8", "Mode", gr.Radio, {"choices": ['INT8', 'INT8_SYM', 'INT4_ASYM', 'INT4_SYM', 'NF4'] if cmd_opts.use_openvino else ['INT8']}), + "nncf_compress_weights": OptionInfo([], "Quantization enabled", gr.CheckboxGroup, {"choices": ["Model", "VAE", "Text Encoder", "ControlNet"], "visible": native}), + "nncf_compress_weights_mode": OptionInfo("INT8", "Quantization type", gr.Dropdown, {"choices": ['INT8', 'INT8_SYM', 'INT4_ASYM', 'INT4_SYM', 'NF4'] if cmd_opts.use_openvino else ['INT8']}), "nncf_compress_weights_raito": OptionInfo(1.0, "Compress ratio", gr.Slider, {"minimum": 0, "maximum": 1, "step": 0.01, "visible": cmd_opts.use_openvino}), "nncf_quantize": OptionInfo([], "OpenVINO enabled", gr.CheckboxGroup, {"choices": ["Model", "VAE", "Text Encoder"], "visible": cmd_opts.use_openvino}), - "nncf_quant_mode": OptionInfo("INT8", "OpenVINO mode", gr.Radio, {"choices": ['INT8', 'FP8_E4M3', 'FP8_E5M2'], "visible": cmd_opts.use_openvino}), + "nncf_quant_mode": OptionInfo("INT8", "OpenVINO mode", gr.Dropdown, {"choices": ['INT8', 'FP8_E4M3', 'FP8_E5M2'], "visible": cmd_opts.use_openvino}), "quant_shuffle_weights": OptionInfo(False, "Shuffle weights", gr.Checkbox, {"visible": native}), }))