Merge branch 'dev' into RUF013

modules/model_quant.py

@@ -1,12 +1,10 @@
 import os
 import re
 import sys
-import copy
 import json
 import time
 import diffusers
-import transformers
-from installer import installed, install, setup_logging
+from installer import install
 from modules.logger import log
 
 
@@ -51,70 +49,6 @@ def dont_quant():
     return False
 
 
-def create_bnb_config(kwargs = None, allow: bool = True, module: str = 'Model', modules_to_not_convert: list | None = None):
-    from modules import shared, devices
-    if allow and (module == 'any' or module in shared.opts.bnb_quantization):
-        load_bnb()
-        if bnb is None:
-            return kwargs
-        bnb_config = diffusers.BitsAndBytesConfig(
-            load_in_8bit=shared.opts.bnb_quantization_type in ['fp8'],
-            load_in_4bit=shared.opts.bnb_quantization_type in ['nf4', 'fp4'],
-            bnb_4bit_quant_storage=shared.opts.bnb_quantization_storage,
-            bnb_4bit_quant_type=shared.opts.bnb_quantization_type,
-            bnb_4bit_compute_dtype=devices.dtype,
-            llm_int8_skip_modules=modules_to_not_convert,
-        )
-        log.debug(f'Quantization: module={module} type=bnb dtype={shared.opts.bnb_quantization_type} storage={shared.opts.bnb_quantization_storage}')
-        if kwargs is None:
-            return bnb_config
-        else:
-            kwargs['quantization_config'] = bnb_config
-            return kwargs
-    return kwargs
-
-
-def create_ao_config(kwargs = None, allow: bool = True, module: str = 'Model', modules_to_not_convert: list | None = None):
-    from modules import shared
-    if allow and (shared.opts.torchao_quantization_mode in {'pre', 'auto'}) and (module == 'any' or module in shared.opts.torchao_quantization):
-        torchao = load_torchao()
-        if torchao is None:
-            return kwargs
-        if module in {'TE', 'LLM'}:
-            ao_config = transformers.TorchAoConfig(quant_type=shared.opts.torchao_quantization_type, modules_to_not_convert=modules_to_not_convert)
-        else:
-            ao_config = diffusers.TorchAoConfig(shared.opts.torchao_quantization_type, modules_to_not_convert=modules_to_not_convert)
-        log.debug(f'Quantization: module={module} type=torchao dtype={shared.opts.torchao_quantization_type}')
-        if kwargs is None:
-            return ao_config
-        else:
-            kwargs['quantization_config'] = ao_config
-            return kwargs
-    return kwargs
-
-
-def create_quanto_config(kwargs = None, allow: bool = True, module: str = 'Model', modules_to_not_convert: list | None = None):
-    from modules import shared
-    if allow and (module == 'any' or module in shared.opts.quanto_quantization):
-        load_quanto(silent=True)
-        if optimum_quanto is None:
-            return kwargs
-        if module in {'TE', 'LLM'}:
-            quanto_config = transformers.QuantoConfig(weights=shared.opts.quanto_quantization_type, modules_to_not_convert=modules_to_not_convert)
-            quanto_config.weights_dtype = quanto_config.weights
-        else:
-            quanto_config = diffusers.QuantoConfig(weights_dtype=shared.opts.quanto_quantization_type, modules_to_not_convert=modules_to_not_convert)
-            quanto_config.activations = None # patch so it works with transformers
-            quanto_config.weights = quanto_config.weights_dtype
-        log.debug(f'Quantization: module={module} type=quanto dtype={shared.opts.quanto_quantization_type}')
-        if kwargs is None:
-            return quanto_config
-        else:
-            kwargs['quantization_config'] = quanto_config
-            return kwargs
-    return kwargs
-
-
 def create_trt_config(kwargs = None, allow: bool = True, module: str = 'Model', modules_to_not_convert: list | None = None):
     from modules import shared
     if allow and (module == 'any' or module in shared.opts.trt_quantization):
@@ -249,7 +183,7 @@ def create_sdnq_config(kwargs = None, allow: bool = True, module: str = 'Model',
 
 def check_quant(module: str = ''):
     from modules import shared
-    if module in shared.opts.sdnq_quantize_weights or module in shared.opts.bnb_quantization or module in shared.opts.torchao_quantization or module in shared.opts.quanto_quantization:
+    if module in shared.opts.sdnq_quantize_weights:
         return True
     return False
 
@@ -286,21 +220,6 @@ def create_config(kwargs = None, allow: bool = True, module: str = 'Model', modu
         if debug:
             log.trace(f'Quantization: type=sdnq config={kwargs.get("quantization_config", None)}')
         return kwargs
-    kwargs = create_bnb_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=bnb config={kwargs.get("quantization_config", None)}')
-        return kwargs
-    kwargs = create_quanto_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=quanto config={kwargs.get("quantization_config", None)}')
-        return kwargs
-    kwargs = create_ao_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=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:
@@ -309,88 +228,6 @@ def create_config(kwargs = None, allow: bool = True, module: str = 'Model', modu
     return kwargs
 
 
-def load_torchao(msg='', silent=False):
-    global ao # pylint: disable=global-statement
-    if ao is not None:
-        return ao
-    if not installed('torchao'):
-        install('torchao==0.10.0', quiet=True)
-        log.warning('Quantization: torchao installed please restart')
-    try:
-        import torchao
-        ao = torchao
-        fn = f'{sys._getframe(2).f_code.co_name}:{sys._getframe(1).f_code.co_name}' # pylint: disable=protected-access
-        log.debug(f'Quantization: type=torchao version={ao.__version__} fn={fn}') # pylint: disable=protected-access
-        from diffusers.utils import import_utils
-        import_utils.is_torchao_available = lambda: True
-        import_utils._torchao_available = True # pylint: disable=protected-access
-        return ao
-    except Exception as e:
-        if len(msg) > 0:
-            log.error(f"{msg} failed to import torchao: {e}")
-        ao = None
-        if not silent:
-            raise
-    return None
-
-
-def load_bnb(msg='', silent=False):
-    from modules import devices
-    global bnb # pylint: disable=global-statement
-    if bnb is not None:
-        return bnb
-    if not installed('bitsandbytes'):
-        if devices.backend == 'cuda':
-            # forcing a version will uninstall the multi-backend-refactor branch of bnb
-            install('bitsandbytes==0.47.0', quiet=True)
-            log.warning('Quantization: bitsandbytes installed please restart')
-    try:
-        import bitsandbytes
-        bnb = bitsandbytes
-        from diffusers.utils import import_utils
-        import_utils._bitsandbytes_available = True # pylint: disable=protected-access
-        import_utils._bitsandbytes_version = '0.43.3' # pylint: disable=protected-access
-        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=bitsandbytes version={bnb.__version__} fn={fn}') # pylint: disable=protected-access
-        return bnb
-    except Exception as e:
-        if len(msg) > 0:
-            log.error(f"{msg} failed to import bitsandbytes: {e}")
-        bnb = None
-        if not silent:
-            raise
-    return None
-
-
-def load_quanto(msg='', silent=False):
-    global optimum_quanto # pylint: disable=global-statement
-    if optimum_quanto is not None:
-        return optimum_quanto
-    if not installed('optimum-quanto'):
-        install('optimum-quanto==0.2.7', quiet=True)
-        log.warning('Quantization: optimum-quanto installed please restart')
-    try:
-        from optimum import quanto # pylint: disable=no-name-in-module
-        # disable device specific tensors because the model can't be moved between cpu and gpu with them
-        quanto.tensor.weights.qbits.WeightQBitsTensor.create = lambda *args, **kwargs: quanto.tensor.weights.qbits.WeightQBitsTensor(*args, **kwargs)
-        optimum_quanto = quanto
-        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=quanto version={quanto.__version__} fn={fn}') # pylint: disable=protected-access
-        from diffusers.utils import import_utils
-        import_utils.is_optimum_quanto_available = lambda: True
-        import_utils._optimum_quanto_available = True # pylint: disable=protected-access
-        import_utils._optimum_quanto_version = quanto.__version__ # pylint: disable=protected-access
-        import_utils._replace_with_quanto_layers = diffusers.quantizers.quanto.utils._replace_with_quanto_layers # pylint: disable=protected-access
-        return optimum_quanto
-    except Exception as e:
-        if len(msg) > 0:
-            log.error(f"{msg} failed to import optimum.quanto: {e}")
-        optimum_quanto = None
-        if not silent:
-            raise
-    return None
-
-
 def load_trt(msg='', silent=False):
     global trt # pylint: disable=global-statement
     if trt is not None:
@@ -642,138 +479,6 @@ def sdnq_quantize_weights(sd_model):
     return sd_model
 
 
-def optimum_quanto_model(model, op=None, sd_model=None, weights=None, activations=None):
-    from modules import devices, shared
-    quanto = load_quanto('Quantize model: type=Optimum Quanto')
-    global quant_last_model_name, quant_last_model_device # pylint: disable=global-statement
-    if model.__class__.__name__ in {"FluxTransformer2DModel", "ChromaTransformer2DModel"}: # LayerNorm is not supported
-        exclude_list = ["transformer_blocks.*.norm1.norm", "transformer_blocks.*.norm2", "transformer_blocks.*.norm1_context.norm", "transformer_blocks.*.norm2_context", "single_transformer_blocks.*.norm.norm", "norm_out.norm"]
-        if model.__class__.__name__ == "ChromaTransformer2DModel":
-            # we ignore the distilled guidance layer because it degrades quality too much
-            # see: https://github.com/huggingface/diffusers/pull/11698#issuecomment-2969717180 for more details
-            exclude_list.append("distilled_guidance_layer.*")
-    elif model.__class__.__name__ == "QwenImageTransformer2DModel":
-        exclude_list = ["transformer_blocks.0.img_mod.1.weight", "time_text_embed", "img_in", "txt_in", "proj_out", "norm_out", "pos_embed"]
-    else:
-        exclude_list = None
-    weights = getattr(quanto, weights) if weights is not None else getattr(quanto, shared.opts.optimum_quanto_weights_type)
-    if activations is not None:
-        activations = getattr(quanto, activations) if activations != 'none' else None
-    elif shared.opts.optimum_quanto_activations_type != 'none':
-        activations = getattr(quanto, shared.opts.optimum_quanto_activations_type)
-    else:
-        activations = None
-    model.eval()
-    backup_embeddings = None
-    if hasattr(model, "get_input_embeddings"):
-        backup_embeddings = copy.deepcopy(model.get_input_embeddings())
-    quanto.quantize(model, weights=weights, activations=activations, exclude=exclude_list)
-    quanto.freeze(model)
-    if hasattr(model, "set_input_embeddings") and backup_embeddings is not None:
-        model.set_input_embeddings(backup_embeddings)
-    if op is not None and shared.opts.optimum_quanto_shuffle_weights:
-        if quant_last_model_name is not None:
-            if "." in quant_last_model_name:
-                last_model_names = quant_last_model_name.split(".")
-                getattr(getattr(sd_model, last_model_names[0]), last_model_names[1]).to(quant_last_model_device)
-            else:
-                getattr(sd_model, quant_last_model_name).to(quant_last_model_device)
-            devices.torch_gc(force=True, reason='quanto')
-        if shared.cmd_opts.medvram or shared.cmd_opts.lowvram or shared.opts.diffusers_offload_mode != "none":
-            quant_last_model_name = op
-            quant_last_model_device = model.device
-        else:
-            quant_last_model_name = None
-            quant_last_model_device = None
-        model.to(devices.device)
-    devices.torch_gc(force=True, reason='quanto')
-    return model
-
-
-def optimum_quanto_weights(sd_model):
-    try:
-        t0 = time.time()
-        from modules import shared, devices, sd_models
-        if shared.opts.diffusers_offload_mode in {"balanced", "sequential"}:
-            log.warning(f"Quantization: type=Optimum.quanto offload={shared.opts.diffusers_offload_mode} not compatible")
-            return sd_model
-        log.info(f"Quantization: type=Optimum.quanto: modules={shared.opts.optimum_quanto_weights}")
-        global quant_last_model_name, quant_last_model_device # pylint: disable=global-statement
-        quanto = load_quanto()
-
-        sd_model = sd_models.apply_function_to_model(sd_model, optimum_quanto_model, shared.opts.optimum_quanto_weights, op="optimum-quanto")
-        if quant_last_model_name is not None:
-            if "." in quant_last_model_name:
-                last_model_names = quant_last_model_name.split(".")
-                getattr(getattr(sd_model, last_model_names[0]), last_model_names[1]).to(quant_last_model_device)
-            else:
-                getattr(sd_model, quant_last_model_name).to(quant_last_model_device)
-            devices.torch_gc(force=True, reason='quanto')
-        quant_last_model_name = None
-        quant_last_model_device = None
-
-        if shared.opts.optimum_quanto_activations_type != 'none':
-            activations = getattr(quanto, shared.opts.optimum_quanto_activations_type)
-        else:
-            activations = None
-
-        if activations is not None:
-            def optimum_quanto_freeze(model, op=None, sd_model=None): # pylint: disable=unused-argument
-                quanto.freeze(model)
-                return model
-            if shared.opts.diffusers_offload_mode == "model":
-                sd_model.enable_model_cpu_offload(device=devices.device)
-                if hasattr(sd_model, "encode_prompt"):
-                    original_encode_prompt = sd_model.encode_prompt
-                    def encode_prompt(*args, **kwargs):
-                        embeds = original_encode_prompt(*args, **kwargs)
-                        sd_model.maybe_free_model_hooks() # Diffusers keeps the TE on VRAM
-                        return embeds
-                    sd_model.encode_prompt = encode_prompt
-            else:
-                sd_models.move_model(sd_model, devices.device)
-            with quanto.Calibration(momentum=0.9):
-                sd_model(prompt="dummy prompt", num_inference_steps=10)
-            sd_model = sd_models.apply_function_to_model(sd_model, optimum_quanto_freeze, shared.opts.optimum_quanto_weights, op="optimum-quanto-freeze")
-            if shared.opts.diffusers_offload_mode == "model":
-                sd_models.disable_offload(sd_model)
-                sd_models.move_model(sd_model, devices.cpu)
-                if hasattr(sd_model, "encode_prompt"):
-                    sd_model.encode_prompt = original_encode_prompt
-            devices.torch_gc(force=True, reason='quanto')
-
-        t1 = time.time()
-        log.info(f"Quantization: type=Optimum.quanto time={t1-t0:.2f}")
-    except Exception as e:
-        log.warning(f"Quantization: type=Optimum.quanto {e}")
-    return sd_model
-
-
-def torchao_quantization(sd_model):
-    from modules import shared, devices, sd_models
-    torchao = load_torchao()
-    q = torchao.quantization
-
-    fn = getattr(q, shared.opts.torchao_quantization_type, None)
-    if fn is None:
-        log.error(f"Quantization: type=TorchAO type={shared.opts.torchao_quantization_type} not supported")
-        return sd_model
-    def torchao_model(model, op=None, sd_model=None): # pylint: disable=unused-argument
-        q.quantize_(model, fn(), device=devices.device)
-        return model
-
-    log.info(f"Quantization: type=TorchAO pipe={sd_model.__class__.__name__} quant={shared.opts.torchao_quantization_type} fn={fn} targets={shared.opts.torchao_quantization}")
-    try:
-        t0 = time.time()
-        sd_models.apply_function_to_model(sd_model, torchao_model, shared.opts.torchao_quantization, op="torchao")
-        t1 = time.time()
-        log.info(f"Quantization: type=TorchAO time={t1-t0:.2f}")
-    except Exception as e:
-        log.error(f"Quantization: type=TorchAO {e}")
-    setup_logging() # torchao uses dynamo which messes with logging so reset is needed
-    return sd_model
-
-
 def get_dit_args(load_config: dict | None = None, module: str | None = None, device_map: bool = False, allow_quant: bool = True, modules_to_not_convert: list | None = None, modules_dtype_dict: dict | None = None):
     from modules import shared, devices
     config = {} if load_config is None else load_config.copy()
@@ -810,12 +515,6 @@ def do_post_load_quant(sd_model, allow=True):
     if shared.opts.sdnq_quantize_weights and (shared.opts.sdnq_quantize_mode == 'post' or (allow and shared.opts.sdnq_quantize_mode == 'auto')):
         log.debug('Load model: post_quant=sdnq')
         sd_model = sdnq_quantize_weights(sd_model)
-    if len(shared.opts.optimum_quanto_weights) > 0:
-        log.debug('Load model: post_quant=quanto')
-        sd_model = optimum_quanto_weights(sd_model)
-    if shared.opts.torchao_quantization and (shared.opts.torchao_quantization_mode == 'post' or (allow and shared.opts.torchao_quantization_mode == 'auto')):
-        log.debug('Load model: post_quant=torchao')
-        sd_model = torchao_quantization(sd_model)
     if shared.opts.layerwise_quantization:
         log.debug('Load model: post_quant=layerwise')
         apply_layerwise(sd_model)