add new optimum-quanto on-the-fly and simplify quantization loading

Signed-off-by: Vladimir Mandic <mandic00@live.com>

CHANGELOG.md

@@ -33,8 +33,10 @@
     against top-10 standard harmful content categories  
   - add banned words/expressions check against prompt variations  
 - **Other**
-  - **upscale**: new [asymmetric vae v2](Heasterian/AsymmetricAutoencoderKLUpscaler_v2) upscaling method
-  - **upscale**: new experimental support for `libvips` upscaling
+  - **upscale**: new [asymmetric vae v2](Heasterian/AsymmetricAutoencoderKLUpscaler_v2) upscaling method  
+  - **upscale**: new experimental support for `libvips` upscaling  
+  - **quantization**: add support for `optimum-quanto` on-the-fly quantization during load for all models  
+    note: previous method for quanto is still valid and is noted in settings as post-load quantization  
   - add remote vae info to metadata, thanks @iDeNoh  
   - add quantization support to **CogView-3Plus**
   - update `diffusers` and other requirements   

cli/prompt-detect.py

@@ -1,32 +0,0 @@
-# Example:
-# > python cli/lang-detect.py "have a good day"
-# > ['eng_latn:1.00']
-# eng=language, latn=latin alphabet, 1.00=confidence
-
-import sys
-import fasttext
-from huggingface_hub import hf_hub_download
-
-
-repo_id = "facebook/fasttext-language-identification"
-model = None
-
-
-def detect(text:str, top:int=1, threshold:float=0.25) -> str:
-    try:
-        global model # pylint: disable=global-statement
-        if model is None:
-            model_path = hf_hub_download(repo_id, filename="model.bin")
-            model = fasttext.load_model(model_path)
-        lang, score = model.predict(text, k=top, threshold=threshold, on_unicode_error="ignore")
-        result = [f"{l.replace("__label__", "").lower()}:{s:.2f}" for l, s in zip(lang, score) if s > threshold][:top]
-        return result
-    except Exception as e:
-        return str(e)
-
-
-if __name__ == "__main__":
-    if len(sys.argv) < 2:
-        print(f"Usage: {sys.argv[0]} <text>")
-    else:
-        print(detect(sys.argv[1]))

modules/model_cogview.py

@@ -18,11 +18,7 @@ def load_common(diffusers_load_config={}, module=None):
     if 'requires_safety_checker' in diffusers_load_config:
         del diffusers_load_config['requires_safety_checker']
 
-    quant_args = {}
-    if not quant_args:
-        quant_args = model_quant.create_bnb_config(quant_args, module=module)
-    if not quant_args:
-        quant_args = model_quant.create_ao_config(quant_args, module=module)
+    quant_args = model_quant.create_config(module=module)
     if quant_args:
         shared.log.debug(f'Load model: type=CogView quantization module="{module}" {quant_args}')
 

modules/model_flux.py

@@ -5,7 +5,7 @@ import diffusers
 import transformers
 from safetensors.torch import load_file
 from huggingface_hub import hf_hub_download
-from modules import shared, devices, modelloader, sd_models, sd_unet, model_te, model_quant
+from modules import shared, errors, devices, modelloader, sd_models, sd_unet, model_te, model_quant
 
 
 debug = shared.log.trace if os.environ.get('SD_LOAD_DEBUG', None) is not None else lambda *args, **kwargs: None
@@ -44,7 +44,6 @@ def load_flux_quanto(checkpoint_info):
     except Exception as e:
         shared.log.error(f"Load model: type=FLUX failed to load Quanto transformer: {e}")
         if debug:
-            from modules import errors
             errors.display(e, 'FLUX Quanto:')
 
     try:
@@ -72,7 +71,6 @@ def load_flux_quanto(checkpoint_info):
     except Exception as e:
         shared.log.error(f"Load model: type=FLUX failed to load Quanto text encoder: {e}")
         if debug:
-            from modules import errors
             errors.display(e, 'FLUX Quanto:')
 
     return transformer, text_encoder_2
@@ -105,33 +103,25 @@ def load_flux_bnb(checkpoint_info, diffusers_load_config): # pylint: disable=unu
         shared.log.error(f"Load model: type=FLUX failed to load BnB transformer: {e}")
         transformer, text_encoder_2 = None, None
         if debug:
-            from modules import errors
             errors.display(e, 'FLUX:')
     return transformer, text_encoder_2
 
 
 def load_quants(kwargs, repo_id, cache_dir, allow_quant):
     try:
-        if not allow_quant:
-            return kwargs
-        quant_args = {}
-        quant_args = model_quant.create_bnb_config(quant_args)
-        if quant_args:
-            model_quant.load_bnb(f'Load model: type=FLUX quant={quant_args}')
-        if not quant_args:
-            quant_args = model_quant.create_ao_config(quant_args)
-            if quant_args:
-                model_quant.load_torchao(f'Load model: type=FLUX quant={quant_args}')
+        quant_args = model_quant.create_config(allow=allow_quant)
         if not quant_args:
             return kwargs
-        if 'transformer' not in kwargs and ('Model' in shared.opts.bnb_quantization or 'Model' in shared.opts.torchao_quantization):
+        if 'transformer' not in kwargs and ('Model' in shared.opts.bnb_quantization or 'Model' in shared.opts.torchao_quantization or 'Model' in shared.opts.quanto_quantization):
             kwargs['transformer'] = diffusers.FluxTransformer2DModel.from_pretrained(repo_id, subfolder="transformer", cache_dir=cache_dir, torch_dtype=devices.dtype, **quant_args)
-            shared.log.debug(f'Quantization: module=transformer type=bnb dtype={shared.opts.bnb_quantization_type} storage={shared.opts.bnb_quantization_storage}')
-        if 'text_encoder_2' not in kwargs and ('Text Encoder' in shared.opts.bnb_quantization or 'Text Encoder' in shared.opts.torchao_quantization):
+        quant_args = model_quant.create_config(allow=allow_quant, module='Text Encoder')
+        if not quant_args:
+            return kwargs
+        if 'text_encoder_2' not in kwargs and ('Text Encoder' in shared.opts.bnb_quantization or 'Text Encoder' in shared.opts.torchao_quantization or 'Text Encoder' in shared.opts.quanto_quantization):
             kwargs['text_encoder_2'] = transformers.T5EncoderModel.from_pretrained(repo_id, subfolder="text_encoder_2", cache_dir=cache_dir, torch_dtype=devices.dtype, **quant_args)
-            shared.log.debug(f'Quantization: module=t5 type=bnb dtype={shared.opts.bnb_quantization_type} storage={shared.opts.bnb_quantization_storage}')
     except Exception as e:
         shared.log.error(f'Quantization: {e}')
+        errors.display(e, 'Quantization:')
     return kwargs
 
 
@@ -197,15 +187,13 @@ def load_transformer(file_path): # triggered by opts.sd_unet change
     else:
         quant_args = model_quant.create_bnb_config({})
         if quant_args:
-            model_quant.load_bnb(f'Load model: type=FLUX quant={quant_args}')
             shared.log.info(f'Load module: type=UNet/Transformer file="{file_path}" offload={shared.opts.diffusers_offload_mode} quant=bnb dtype={devices.dtype}')
             from modules.model_flux_nf4 import load_flux_nf4
             transformer, _text_encoder_2 = load_flux_nf4(file_path, prequantized=False)
             if transformer is not None:
                 return transformer
-        quant_args = model_quant.create_ao_config({})
+        quant_args = model_quant.create_config()
         if quant_args:
-            model_quant.load_torchao(f'Load model: type=FLUX quant={quant_args}')
             shared.log.info(f'Load module: type=UNet/Transformer file="{file_path}" offload={shared.opts.diffusers_offload_mode} quant=torchao dtype={devices.dtype}')
             transformer = diffusers.FluxTransformer2DModel.from_single_file(file_path, **diffusers_load_config, **quant_args)
             if transformer is not None:
@@ -249,7 +237,6 @@ def load_flux(checkpoint_info, diffusers_load_config): # triggered by opts.sd_ch
             shared.log.error(f"Load model: type=FLUX failed to load UNet: {e}")
             shared.opts.sd_unet = 'Default'
             if debug:
-                from modules import errors
                 errors.display(e, 'FLUX UNet:')
     if shared.opts.sd_text_encoder != 'Default':
         try:
@@ -263,7 +250,6 @@ def load_flux(checkpoint_info, diffusers_load_config): # triggered by opts.sd_ch
             shared.log.error(f"Load model: type=FLUX failed to load T5: {e}")
             shared.opts.sd_text_encoder = 'Default'
             if debug:
-                from modules import errors
                 errors.display(e, 'FLUX T5:')
     if shared.opts.sd_vae != 'Default' and shared.opts.sd_vae != 'Automatic':
         try:
@@ -278,7 +264,6 @@ def load_flux(checkpoint_info, diffusers_load_config): # triggered by opts.sd_ch
             shared.log.error(f"Load model: type=FLUX failed to load VAE: {e}")
             shared.opts.sd_vae = 'Default'
             if debug:
-                from modules import errors
                 errors.display(e, 'FLUX VAE:')
 
     # load quantized components if any
@@ -293,7 +278,6 @@ def load_flux(checkpoint_info, diffusers_load_config): # triggered by opts.sd_ch
         except Exception as e:
             shared.log.error(f"Load model: type=FLUX failed to load NF4 components: {e}")
             if debug:
-                from modules import errors
                 errors.display(e, 'FLUX NF4:')
     if quant == 'qint8' or quant == 'qint4':
         try:
@@ -305,7 +289,6 @@ def load_flux(checkpoint_info, diffusers_load_config): # triggered by opts.sd_ch
         except Exception as e:
             shared.log.error(f"Load model: type=FLUX failed to load Quanto components: {e}")
             if debug:
-                from modules import errors
                 errors.display(e, 'FLUX Quanto:')
 
     # initialize pipeline with pre-loaded components
@@ -346,8 +329,7 @@ def load_flux(checkpoint_info, diffusers_load_config): # triggered by opts.sd_ch
     fn = checkpoint_info.path
     if (fn is None) or (not os.path.exists(fn) or os.path.isdir(fn)):
         kwargs = load_quants(kwargs, repo_id, cache_dir=shared.opts.diffusers_dir, allow_quant=allow_quant)
-    kwargs = model_quant.create_bnb_config(kwargs, allow_quant)
-    kwargs = model_quant.create_ao_config(kwargs, allow_quant)
+    # kwargs = model_quant.create_config(kwargs, allow_quant)
     if fn.endswith('.safetensors') and os.path.isfile(fn):
         pipe = diffusers.FluxPipeline.from_single_file(fn, cache_dir=shared.opts.diffusers_dir, **kwargs, **diffusers_load_config)
     else:

modules/model_lumina.py

@@ -32,9 +32,7 @@ def load_lumina2(checkpoint_info, diffusers_load_config={}):
     if quant_args:
         model_quant.load_bnb(f'Load model: type=Lumina quant={quant_args}')
     if not quant_args:
-        quant_args = model_quant.create_ao_config(quant_args)
-        if quant_args:
-            model_quant.load_torchao(f'Load model: type=Lumina quant={quant_args}')
+        quant_args = model_quant.create_config()
     kwargs = {}
     repo_id = sd_models.path_to_repo(checkpoint_info.name)
     if ('Model' in shared.opts.bnb_quantization or 'Model' in shared.opts.torchao_quantization):

modules/model_quant.py

@@ -1,3 +1,4 @@
+import os
 import sys
 import copy
 import time
@@ -9,9 +10,9 @@ ao = None
 bnb = None
 intel_nncf = None
 optimum_quanto = None
-
 quant_last_model_name = None
 quant_last_model_device = None
+debug = os.environ.get('SD_QUANT_DEBUG', None) is not None
 
 
 def get_quant(name):
@@ -44,7 +45,7 @@ def create_bnb_config(kwargs = None, allow_bnb: bool = True, module: str = 'Mode
                 bnb_4bit_quant_type=shared.opts.bnb_quantization_type,
                 bnb_4bit_compute_dtype=devices.dtype
             )
-            shared.log.debug(f'Quantization: module=all type=bnb dtype={shared.opts.bnb_quantization_type} storage={shared.opts.bnb_quantization_storage}')
+            log.debug(f'Quantization: module=all type=bnb dtype={shared.opts.bnb_quantization_type} storage={shared.opts.bnb_quantization_storage}')
             if kwargs is None:
                 return bnb_config
             else:
@@ -60,9 +61,8 @@ def create_ao_config(kwargs = None, allow_ao: bool = True, module: str = 'Model'
             load_torchao()
             if ao is None:
                 return kwargs
-            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}')
+            log.debug(f'Quantization: module=all type=torchao dtype={shared.opts.torchao_quantization_type}')
             if kwargs is None:
                 return ao_config
             else:
@@ -71,6 +71,47 @@ def create_ao_config(kwargs = None, allow_ao: bool = True, module: str = 'Model'
     return kwargs
 
 
+def create_quanto_config(kwargs = None, allow_quanto: bool = True, module: str = 'Model'):
+    from modules import shared
+    if len(shared.opts.quanto_quantization) > 0 and allow_quanto:
+        if 'Model' in shared.opts.quanto_quantization or (module is not None and module in shared.opts.quanto_quantization):
+            load_quanto(silent=True)
+            if optimum_quanto is None:
+                return kwargs
+            quanto_config = diffusers.QuantoConfig(
+                weights_dtype=shared.opts.quanto_quantization_type,
+            )
+            quanto_config.activations = None # patch so it works with transformers
+            log.debug(f'Quantization: module=all 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_config(kwargs = None, allow: bool = True, module: str = 'Model'):
+    if kwargs is None:
+        kwargs = {}
+    kwargs = create_bnb_config(kwargs, allow_bnb=allow, module=module)
+    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_ao_config(kwargs, allow_ao=allow, module=module)
+    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_quanto_config(kwargs, allow_quanto=allow, module=module)
+    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
+    return kwargs
+
+
 def load_torchao(msg='', silent=False):
     global ao # pylint: disable=global-statement
     if ao is not None:
@@ -81,6 +122,9 @@ def load_torchao(msg='', silent=False):
         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:
@@ -102,9 +146,10 @@ def load_bnb(msg='', silent=False):
     try:
         import bitsandbytes
         bnb = bitsandbytes
-        diffusers.utils.import_utils._bitsandbytes_available = True # pylint: disable=protected-access
-        diffusers.utils.import_utils._bitsandbytes_version = '0.43.3' # pylint: disable=protected-access
-        fn = f'{sys._getframe(2).f_code.co_name}:{sys._getframe(1).f_code.co_name}' # pylint: disable=protected-access
+        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:
@@ -117,18 +162,20 @@ def load_bnb(msg='', silent=False):
 
 
 def load_quanto(msg='', silent=False):
-    from modules import shared
     global optimum_quanto # pylint: disable=global-statement
     if optimum_quanto is not None:
         return optimum_quanto
-    install('optimum-quanto==0.2.6', quiet=True)
+    install('optimum-quanto==0.2.7', quiet=True)
     try:
         from optimum import quanto # pylint: disable=no-name-in-module
         optimum_quanto = quanto
-        fn = f'{sys._getframe(2).f_code.co_name}:{sys._getframe(1).f_code.co_name}' # 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=quanto version={quanto.__version__} fn={fn}') # pylint: disable=protected-access
-        if shared.opts.diffusers_offload_mode in {'balanced', 'sequential'}:
-            shared.log.error(f'Quantization: type=quanto offload={shared.opts.diffusers_offload_mode} not supported')
+        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:
@@ -169,7 +216,7 @@ def apply_layerwise(sd_model, quiet:bool=False):
         storage_dtype = torch.float8_e5m2
     else:
         storage_dtype = None
-        shared.log.warning(f'Quantization: type=layerwise storage={shared.opts.layerwise_quantization_storage} not supported')
+        log.warning(f'Quantization: type=layerwise storage={shared.opts.layerwise_quantization_storage} not supported')
         return
     non_blocking = False
     if not hasattr(quantization_config.QuantizationMethod, 'LAYERWISE'):
@@ -198,7 +245,7 @@ def apply_layerwise(sd_model, quiet:bool=False):
                     m.quantization_method = quantization_config.QuantizationMethod.LAYERWISE # pylint: disable=no-member
                     log.quiet(quiet, f'Quantization: type=layerwise module={module} cls={cls} storage={storage_dtype} compute={devices.dtype} blocking={not non_blocking}')
         except Exception as e:
-            shared.log.error(f'Quantization: type=layerwise {e}')
+            log.error(f'Quantization: type=layerwise {e}')
 
 
 def nncf_send_to_device(model, device):
@@ -244,7 +291,7 @@ def nncf_compress_weights(sd_model):
     try:
         t0 = time.time()
         from modules import shared, devices, sd_models
-        shared.log.info(f"Quantization: type=NNCF modules={shared.opts.nncf_compress_weights}")
+        log.info(f"Quantization: type=NNCF modules={shared.opts.nncf_compress_weights}")
         global quant_last_model_name, quant_last_model_device # pylint: disable=global-statement
 
         sd_model = sd_models.apply_function_to_model(sd_model, nncf_compress_model, shared.opts.nncf_compress_weights, op="nncf")
@@ -259,9 +306,9 @@ def nncf_compress_weights(sd_model):
         quant_last_model_device = None
 
         t1 = time.time()
-        shared.log.info(f"Quantization: type=NNCF time={t1-t0:.2f}")
+        log.info(f"Quantization: type=NNCF time={t1-t0:.2f}")
     except Exception as e:
-        shared.log.warning(f"Quantization: type=NNCF {e}")
+        log.warning(f"Quantization: type=NNCF {e}")
     return sd_model
 
 
@@ -312,9 +359,9 @@ def optimum_quanto_weights(sd_model):
         t0 = time.time()
         from modules import shared, devices, sd_models
         if shared.opts.diffusers_offload_mode in {"balanced", "sequential"}:
-            shared.log.warning(f"Quantization: type=Optimum.quanto offload={shared.opts.diffusers_offload_mode} not compatible")
+            log.warning(f"Quantization: type=Optimum.quanto offload={shared.opts.diffusers_offload_mode} not compatible")
             return sd_model
-        shared.log.info(f"Quantization: type=Optimum.quanto: modules={shared.opts.optimum_quanto_weights}")
+        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()
         quanto.tensor.qbits.QBitsTensor.create = lambda *args, **kwargs: quanto.tensor.qbits.QBitsTensor(*args, **kwargs)
@@ -361,9 +408,9 @@ def optimum_quanto_weights(sd_model):
             devices.torch_gc(force=True)
 
         t1 = time.time()
-        shared.log.info(f"Quantization: type=Optimum.quanto time={t1-t0:.2f}")
+        log.info(f"Quantization: type=Optimum.quanto time={t1-t0:.2f}")
     except Exception as e:
-        shared.log.warning(f"Quantization: type=Optimum.quanto {e}")
+        log.warning(f"Quantization: type=Optimum.quanto {e}")
     return sd_model
 
 
@@ -374,19 +421,19 @@ def torchao_quantization(sd_model):
 
     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")
+        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
 
-    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}")
+    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()
-        shared.log.info(f"Quantization: type=TorchAO time={t1-t0:.2f}")
+        log.info(f"Quantization: type=TorchAO time={t1-t0:.2f}")
     except Exception as e:
-        shared.log.error(f"Quantization: type=TorchAO {e}")
+        log.error(f"Quantization: type=TorchAO {e}")
     setup_logging() # torchao uses dynamo which messes with logging so reset is needed
     return sd_model

modules/model_sana.py

@@ -8,13 +8,7 @@ from modules import shared, sd_models, devices, modelloader, model_quant
 
 def load_quants(kwargs, repo_id, cache_dir):
     quant_args = {}
-    quant_args = model_quant.create_bnb_config(quant_args)
-    if quant_args:
-        model_quant.load_bnb(f'Load model: type=Sana quant={quant_args}')
-    if not quant_args:
-        quant_args = model_quant.create_ao_config(quant_args)
-        if quant_args:
-            model_quant.load_torchao(f'Load model: type=Sana quant={quant_args}')
+    quant_args = model_quant.create_config()
     if not quant_args:
         return kwargs
     load_args = kwargs.copy()

modules/model_sd3.py

@@ -52,14 +52,7 @@ def load_overrides(kwargs, cache_dir):
 
 
 def load_quants(kwargs, repo_id, cache_dir):
-    quant_args = {}
-    quant_args = model_quant.create_bnb_config(quant_args)
-    if quant_args:
-        model_quant.load_bnb(f'Load model: type=SD3 quant={quant_args}')
-    if not quant_args:
-        quant_args = model_quant.create_ao_config(quant_args)
-        if quant_args:
-            model_quant.load_torchao(f'Load model: type=SD3 quant={quant_args}')
+    quant_args = model_quant.create_config()
     if not quant_args:
         return kwargs
     if 'Model' in shared.opts.bnb_quantization and 'transformer' not in kwargs:
@@ -157,8 +150,7 @@ def load_sd3(checkpoint_info, cache_dir=None, config=None):
 
     shared.log.debug(f'Load model: type=SD3 kwargs={list(kwargs)} repo="{repo_id}"')
 
-    kwargs = model_quant.create_bnb_config(kwargs)
-    kwargs = model_quant.create_ao_config(kwargs)
+    kwargs = model_quant.create_config(kwargs)
     pipe = loader(
         repo_id,
         torch_dtype=devices.dtype,

modules/model_tools.py

@@ -69,13 +69,11 @@ def load_modules(repo_id: str, params: dict):
             subfolder = 'text_encoder_2'
         if cls == transformers.T5EncoderModel: # t5-xxl
             subfolder = 'text_encoder_3'
-            kwargs = model_quant.create_bnb_config(kwargs)
-            kwargs = model_quant.create_ao_config(kwargs)
+            kwargs = model_quant.create_config(kwargs)
             kwargs['variant'] = 'fp16'
         if cls == diffusers.SD3Transformer2DModel:
             subfolder = 'transformer'
-            kwargs = model_quant.create_bnb_config(kwargs)
-            kwargs = model_quant.create_ao_config(kwargs)
+            kwargs = model_quant.create_config(kwargs)
         if subfolder is None:
             continue
         shared.log.debug(f'Load: module={name} class={cls.__name__} repo={repo_id} location={subfolder}')

modules/shared.py

@@ -513,7 +513,11 @@ options_templates.update(options_section(('quantization', "Quantization Settings
     "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("<h2>Optimum Quanto</h2>", "", gr.HTML),
+    "quanto_quantization_sep": OptionInfo("<h2>Optimum Quanto</h2>", "", gr.HTML),
+    "quanto_quantization": OptionInfo([], "Quantization enabled", gr.CheckboxGroup, {"choices": ["Model", "VAE", "Text Encoder"], "visible": native}),
+    "quanto_quantization_type": OptionInfo("int8", "Quantization weights type", gr.Dropdown, {"choices": ["float8", "int8", "int4", "int2"], "visible": native}),
+
+    "optimum_quanto_sep": OptionInfo("<h2>Optimum Quanto: post-load</h2>", "", gr.HTML),
     "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}),
@@ -524,7 +528,7 @@ options_templates.update(options_section(('quantization', "Quantization Settings
     "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_compress_sep": OptionInfo("<h2>NNCF</h2>", "", gr.HTML),
+    "nncf_compress_sep": OptionInfo("<h2>NNCF: Neural Network Compression Framework</h2>", "", gr.HTML),
     "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(0, "Compress ratio", gr.Slider, {"minimum": 0, "maximum": 1, "step": 0.01, "visible": cmd_opts.use_openvino}),

scripts/allegrovideo.py

@@ -61,14 +61,7 @@ class Script(scripts.Script):
         if shared.sd_model.__class__ != diffusers.AllegroPipeline:
             sd_models.unload_model_weights()
             t0 = time.time()
-            quant_args = {}
-            quant_args = model_quant.create_bnb_config(quant_args)
-            if quant_args:
-                model_quant.load_bnb(f'Load model: type=Allegro quant={quant_args}')
-            if not quant_args:
-                quant_args = model_quant.create_ao_config(quant_args)
-                if quant_args:
-                    model_quant.load_torchao(f'Load model: type=Allegro quant={quant_args}')
+            quant_args = model_quant.create_config()
             transformer = diffusers.AllegroTransformer3DModel.from_pretrained(
                 repo_id,
                 subfolder="transformer",

scripts/hunyuanvideo.py

@@ -91,14 +91,7 @@ class Script(scripts.Script):
         if shared.sd_model.__class__ != diffusers.HunyuanVideoPipeline or model != loaded_model:
             sd_models.unload_model_weights()
             t0 = time.time()
-            quant_args = {}
-            quant_args = model_quant.create_bnb_config(quant_args)
-            if quant_args:
-                model_quant.load_bnb(f'Load model: type=HunyuanVideo quant={quant_args}')
-            if not quant_args:
-                quant_args = model_quant.create_ao_config(quant_args)
-                if quant_args:
-                    model_quant.load_torchao(f'Load model: type=HunyuanVideo quant={quant_args}')
+            quant_args = model_quant.create_config()
             transformer = diffusers.HunyuanVideoTransformer3DModel.from_pretrained(
                 pretrained_model_name_or_path='tencent/HunyuanVideo',
                 subfolder="transformer",

scripts/ltxvideo.py

@@ -16,14 +16,7 @@ repos = {
 
 
 def load_quants(kwargs, repo_id):
-    quant_args = {}
-    quant_args = model_quant.create_bnb_config(quant_args)
-    if quant_args:
-        model_quant.load_bnb(f'Load model: type=LTXVideo quant={quant_args}')
-    if not quant_args:
-        quant_args = model_quant.create_ao_config(quant_args)
-        if quant_args:
-            model_quant.load_torchao(f'Load model: type=LTXVideo quant={quant_args}')
+    quant_args = model_quant.create_config()
     if not quant_args:
         return kwargs
     model_quant.load_bnb(f'Load model: type=LTX quant={quant_args}')
@@ -119,9 +112,7 @@ class Script(scripts.Script):
             repo_id = model_custom
         if shared.sd_model.__class__ != cls:
             sd_models.unload_model_weights()
-            kwargs = {}
-            kwargs = model_quant.create_bnb_config(kwargs)
-            kwargs = model_quant.create_ao_config(kwargs)
+            kwargs = model_quant.create_config()
             diffusers.LTXVideoTransformer3DModel.forward = teacache_forward
             if os.path.isfile(repo_id):
                 shared.sd_model = cls.from_single_file(

scripts/mochivideo.py

@@ -42,9 +42,7 @@ class Script(scripts.Script):
         cls = diffusers.MochiPipeline
         if shared.sd_model.__class__ != cls:
             sd_models.unload_model_weights()
-            kwargs = {}
-            kwargs = model_quant.create_bnb_config(kwargs)
-            kwargs = model_quant.create_ao_config(kwargs)
+            kwargs = model_quant.create_config()
             shared.sd_model = cls.from_pretrained(
                 repo_id,
                 cache_dir = shared.opts.hfcache_dir,