add torchao

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

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  

TODO.md

@@ -10,21 +10,14 @@ Main ToDo list can be found at [GitHub projects](https://github.com/users/vladma
 
 ## Future Candidates
 
-- SD35 IPAdapter: <https://github.com/huggingface/diffusers/issues/9966>
-- SD35 LoRA: <https://github.com/huggingface/diffusers/issues/9950>
-- Flux IPAdapter: <https://github.com/huggingface/diffusers/issues/9825>
+- SD35 IPAdapter: <https://github.com/huggingface/diffusers/pull/9987>
+- Flux IPAdapter: <https://github.com/huggingface/diffusers/pull/10261>
 - Flux NF4: <https://github.com/huggingface/diffusers/issues/9996>
-- SANA: <https://github.com/huggingface/diffusers/pull/9982>
 - LTX-Video: <https://github.com/huggingface/diffusers/pull/10021> <https://huggingface.co/Lightricks/LTX-Video> <https://huggingface.co/spaces/Lightricks/LTX-Video-Playground/tree/main>
-- TorchAO: <https://github.com/huggingface/diffusers/pull/10009>
+- GGUF: <https://github.com/huggingface/diffusers/pull/9964>
 
 ## Other
 
 - IPAdapter negative: <https://github.com/huggingface/diffusers/discussions/7167>
 - Control API enhance scripts compatibility
 - PixelSmith: <https://github.com/Thanos-DB/Pixelsmith>
-
-## Workaround in place
-
-- GGUF <https://github.com/huggingface/diffusers/issues/9487>
-- FlowMatch <https://github.com/huggingface/diffusers/issues/9607> <https://github.com/huggingface/diffusers/issues/9924>

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():

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:

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

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)

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}')

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'):

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}")

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("<h2>BitsAndBytes</h2>", "", 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("<h2>Optimum Quanto</h2>", "", 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("<h2>TorchAO</h2>", "", 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("<h2>NNCF</h2>", "", 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}),
 }))