Optimum Quanto activations support

modules/model_t5.py

@@ -49,7 +49,7 @@ def load_t5(t5=None, cache_dir=None):
             cache_dir=cache_dir,
             torch_dtype=devices.dtype,
         )
-        t5 = optimum_quanto_model(t5, weights="qint8")
+        t5 = optimum_quanto_model(t5, weights="qint8", activations="none")
     elif 'int8' in t5.lower():
         modelloader.hf_login()
         from installer import install

modules/sd_models_compile.py

@@ -185,16 +185,23 @@ def nncf_compress_weights(sd_model):
         shared.log.warning(f"NNCF Compress Weights: error: {e}")
     return sd_model
 
-def optimum_quanto_model(model, op=None, sd_model=None, weights=None):
+def optimum_quanto_model(model, op=None, sd_model=None, weights=None, activations=None):
     from optimum import quanto
     global quant_last_model_name, quant_last_model_device
     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)
-    quanto.freeze(model)
+    quanto.quantize(model, weights=weights, activations=activations)
+    if activations is not None:
+        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.quant_shuffle_weights:
@@ -239,6 +246,27 @@ def optimum_quanto_weights(sd_model):
         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):
+                quanto.freeze(model)
+                return model
+            if shared.opts.diffusers_offload_mode == "model":
+                sd_model.enable_model_cpu_offload(device=devices.device)
+            else:
+                sd_models.move_model(sd_model, devices.device)
+            with quanto.Calibration(momentum=0.9):
+                sd_model(prompt="dummy", height=512, width=512, guidance_scale=4.0, num_inference_steps=10)
+            sd_model = apply_compile_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)
+            devices.torch_gc(force=True)
+
         t1 = time.time()
         shared.log.info(f"Optimum Quanto Weights: time={t1-t0:.2f}")
     except Exception as e:

modules/shared.py

@@ -469,6 +469,7 @@ options_templates.update(options_section(('cuda', "Compute Settings"), {
     "nncf_compress_weights": OptionInfo([], "Compress Model weights with NNCF INT8", gr.CheckboxGroup, {"choices": ["Model", "VAE", "Text Encoder", "ControlNet"], "visible": native}),
     "optimum_quanto_weights": OptionInfo([], "Quantize Model weights with Optimum Quanto", gr.CheckboxGroup, {"choices": ["Model", "VAE", "Text Encoder", "ControlNet"], "visible": native}),
     "optimum_quanto_weights_type": OptionInfo("qint8", "Quant mode for Optimum Quanto", gr.Radio, {"choices": ['qint8', 'qfloat8_e4m3fn', 'qfloat8_e5m2', 'qint4', 'qint2'], "visible": native}),
+    "optimum_quanto_activations_type": OptionInfo("none", "Activation Quant mode for Optimum Quanto", gr.Radio, {"choices": ['none', 'qint8', 'qfloat8_e4m3fn', 'qfloat8_e5m2'], "visible": native}),
 
     "ipex_sep": OptionInfo("<h2>IPEX</h2>", "", gr.HTML, {"visible": devices.backend == "ipex"}),
     "ipex_optimize": OptionInfo([], "IPEX Optimize for Intel GPUs", gr.CheckboxGroup, {"choices": ["Model", "VAE", "Text Encoder", "Upscaler"], "visible": devices.backend == "ipex"}),