From dc9e60aa671cbbc129e6271c939a8ead7b62ea3f Mon Sep 17 00:00:00 2001 From: Disty0 Date: Sun, 4 Aug 2024 04:46:06 +0300 Subject: [PATCH] Quant add shuffle models option --- modules/sd_models_compile.py | 94 +++++++++++++++++++++++++++--------- modules/shared.py | 1 + 2 files changed, 71 insertions(+), 24 deletions(-) diff --git a/modules/sd_models_compile.py b/modules/sd_models_compile.py index 2c516b620..67402fc9e 100644 --- a/modules/sd_models_compile.py +++ b/modules/sd_models_compile.py @@ -24,39 +24,30 @@ class CompiledModelState: self.partitioned_modules = {} +quant_last_model_name = None +quant_last_model_device = None deepcache_worker = None def apply_compile_to_model(sd_model, function, options, op=None): if "Model" in options: if hasattr(sd_model, 'unet') and hasattr(sd_model.unet, 'config'): - sd_model.unet = function(sd_model.unet) + sd_model.unet = function(sd_model.unet, op="unet", sd_model=sd_model) if hasattr(sd_model, 'transformer') and hasattr(sd_model.transformer, 'config'): - sd_model.transformer = function(sd_model.transformer) + sd_model.transformer = function(sd_model.transformer, op="transformer", sd_model=sd_model) if hasattr(sd_model, 'decoder_pipe') and hasattr(sd_model, 'decoder'): sd_model.decoder = None - sd_model.decoder = sd_model.decoder_pipe.decoder = function(sd_model.decoder_pipe.decoder) + sd_model.decoder = sd_model.decoder_pipe.decoder = function(sd_model.decoder_pipe.decoder, op="decoder_pipe.decoder", sd_model=sd_model) if hasattr(sd_model, 'prior_pipe') and hasattr(sd_model.prior_pipe, 'prior'): if op == "nncf" and "StableCascade" in sd_model.__class__.__name__: # fixes dtype errors backup_clip_txt_pooled_mapper = copy.deepcopy(sd_model.prior_pipe.prior.clip_txt_pooled_mapper) - sd_model.prior_pipe.prior = function(sd_model.prior_pipe.prior) + sd_model.prior_pipe.prior = function(sd_model.prior_pipe.prior, op="prior_pipe.prior", sd_model=sd_model) if op == "nncf" and "StableCascade" in sd_model.__class__.__name__: sd_model.prior_pipe.prior.clip_txt_pooled_mapper = backup_clip_txt_pooled_mapper - if "VAE" in options: - if hasattr(sd_model, 'vae') and hasattr(sd_model.vae, 'decode'): - sd_model.vae = function(sd_model.vae) - if hasattr(sd_model, 'movq') and hasattr(sd_model.movq, 'decode'): - sd_model.movq = function(sd_model.movq) - if hasattr(sd_model, 'vqgan') and hasattr(sd_model.vqgan, 'decode'): - sd_model.vqgan = function(sd_model.vqgan) - if hasattr(sd_model, 'decoder_pipe') and hasattr(sd_model.decoder_pipe, 'vqgan'): - sd_model.decoder_pipe.vqgan = sd_model.vqgan - if hasattr(sd_model, 'image_encoder') and hasattr(sd_model.image_encoder, 'config'): - sd_model.image_encoder = function(sd_model.image_encoder) if "Text Encoder" in options: if hasattr(sd_model, 'text_encoder') and hasattr(sd_model.text_encoder, 'config'): if hasattr(sd_model, 'decoder_pipe') and hasattr(sd_model.decoder_pipe, 'text_encoder'): - sd_model.decoder_pipe.text_encoder = function(sd_model.decoder_pipe.text_encoder) + sd_model.decoder_pipe.text_encoder = function(sd_model.decoder_pipe.text_encoder, op="decoder_pipe.text_encoder", sd_model=sd_model) else: if op == "nncf" and sd_model.text_encoder.__class__.__name__ in {"T5EncoderModel", "UMT5EncoderModel"}: from modules.sd_hijack import NNCF_T5DenseGatedActDense # T5DenseGatedActDense uses fp32 @@ -65,7 +56,7 @@ def apply_compile_to_model(sd_model, function, options, op=None): sd_model.text_encoder.encoder.block[i].layer[1].DenseReluDense, dtype=torch.float32 if devices.dtype != torch.bfloat16 else torch.bfloat16 ) - sd_model.text_encoder = function(sd_model.text_encoder) + sd_model.text_encoder = function(sd_model.text_encoder, op="text_encoder", sd_model=sd_model) if hasattr(sd_model, 'text_encoder_2') and hasattr(sd_model.text_encoder_2, 'config'): if op == "nncf" and sd_model.text_encoder_2.__class__.__name__ in {"T5EncoderModel", "UMT5EncoderModel"}: from modules.sd_hijack import NNCF_T5DenseGatedActDense # T5DenseGatedActDense uses fp32 @@ -74,7 +65,7 @@ def apply_compile_to_model(sd_model, function, options, op=None): sd_model.text_encoder_2.encoder.block[i].layer[1].DenseReluDense, dtype=torch.float32 if devices.dtype != torch.bfloat16 else torch.bfloat16 ) - sd_model.text_encoder_2 = function(sd_model.text_encoder_2) + sd_model.text_encoder_2 = function(sd_model.text_encoder_2, op="text_encoder_2", sd_model=sd_model) if hasattr(sd_model, 'text_encoder_3') and hasattr(sd_model.text_encoder_3, 'config'): if op == "nncf" and sd_model.text_encoder_3.__class__.__name__ in {"T5EncoderModel", "UMT5EncoderModel"}: from modules.sd_hijack import NNCF_T5DenseGatedActDense # T5DenseGatedActDense uses fp32 @@ -83,9 +74,20 @@ def apply_compile_to_model(sd_model, function, options, op=None): sd_model.text_encoder_3.encoder.block[i].layer[1].DenseReluDense, dtype=torch.float32 if devices.dtype != torch.bfloat16 else torch.bfloat16 ) - sd_model.text_encoder_3 = function(sd_model.text_encoder_3) + sd_model.text_encoder_3 = function(sd_model.text_encoder_3, op="text_encoder_3", sd_model=sd_model) if hasattr(sd_model, 'prior_pipe') and hasattr(sd_model.prior_pipe, 'text_encoder'): - sd_model.prior_pipe.text_encoder = function(sd_model.prior_pipe.text_encoder) + sd_model.prior_pipe.text_encoder = function(sd_model.prior_pipe.text_encoder, op="prior_pipe.text_encoder", sd_model=sd_model) + if "VAE" in options: + if hasattr(sd_model, 'vae') and hasattr(sd_model.vae, 'decode'): + sd_model.vae = function(sd_model.vae, op="vae", sd_model=sd_model) + if hasattr(sd_model, 'movq') and hasattr(sd_model.movq, 'decode'): + sd_model.movq = function(sd_model.movq, op="movq", sd_model=sd_model) + if hasattr(sd_model, 'vqgan') and hasattr(sd_model.vqgan, 'decode'): + sd_model.vqgan = function(sd_model.vqgan, op="vqgan", sd_model=sd_model) + if hasattr(sd_model, 'decoder_pipe') and hasattr(sd_model.decoder_pipe, 'vqgan'): + sd_model.decoder_pipe.vqgan = sd_model.vqgan + if hasattr(sd_model, 'image_encoder') and hasattr(sd_model.image_encoder, 'config'): + sd_model.image_encoder = function(sd_model.image_encoder, op="image_encoder", sd_model=sd_model) return sd_model @@ -94,7 +96,7 @@ def ipex_optimize(sd_model): try: t0 = time.time() - def ipex_optimize_model(model): + def ipex_optimize_model(model, op=None, sd_model=None): import intel_extension_for_pytorch as ipex # pylint: disable=import-error, unused-import model.eval() model.training = False @@ -129,8 +131,9 @@ def nncf_send_to_device(model): child.zero_point = child.zero_point.to(devices.device) nncf_send_to_device(child) -def nncf_compress_model(model): +def nncf_compress_model(model, op=None, sd_model=None): import nncf + global quant_last_model_name, quant_last_model_device model.eval() backup_embeddings = None if hasattr(model, "get_input_embeddings"): @@ -139,6 +142,17 @@ def nncf_compress_model(model): nncf_send_to_device(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: + 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) + quant_last_model_name = op + quant_last_model_device = model.device + model.to(devices.device) devices.torch_gc(force=True) return model @@ -146,10 +160,20 @@ def nncf_compress_weights(sd_model): try: t0 = time.time() shared.log.info(f"NNCF Compress Weights: {shared.opts.nncf_compress_weights}") + global quant_last_model_name, quant_last_model_device from installer import install install('nncf==2.7.0', quiet=True) sd_model = apply_compile_to_model(sd_model, nncf_compress_model, shared.opts.nncf_compress_weights, op="nncf") + 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) + quant_last_model_name = None + quant_last_model_device = None t1 = time.time() shared.log.info(f"NNCF Compress Weights: time={t1-t0:.2f}") @@ -157,8 +181,9 @@ def nncf_compress_weights(sd_model): shared.log.warning(f"NNCF Compress Weights: error: {e}") return sd_model -def optimum_quanto_model(model, weights=None): +def optimum_quanto_model(model, op=None, sd_model=None, weights=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) model.eval() backup_embeddings = None @@ -168,6 +193,17 @@ def optimum_quanto_model(model, weights=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: + 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) + quant_last_model_name = op + quant_last_model_device = model.device + model.to(devices.device) devices.torch_gc(force=True) return model @@ -175,10 +211,20 @@ def optimum_quanto_weights(sd_model): try: t0 = time.time() shared.log.info(f"Optimum Quanto Weights: {shared.opts.optimum_quanto_weights}") + global quant_last_model_name, quant_last_model_device from installer import install install('optimum-quanto', quiet=True) sd_model = apply_compile_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) + quant_last_model_name = None + quant_last_model_device = None t1 = time.time() shared.log.info(f"Optimum Quanto Weights: time={t1-t0:.2f}") @@ -283,7 +329,7 @@ def compile_torch(sd_model): torch._dynamo.reset() # pylint: disable=protected-access shared.log.debug(f"Model compile available backends: {torch._dynamo.list_backends()}") # pylint: disable=protected-access - def torch_compile_model(model): + def torch_compile_model(model, op=None, sd_model=None): if model.device.type != "meta": return_device = model.device model = torch.compile(model.to(devices.device), diff --git a/modules/shared.py b/modules/shared.py index 784e90524..b5ab3fac8 100644 --- a/modules/shared.py +++ b/modules/shared.py @@ -449,6 +449,7 @@ options_templates.update(options_section(('cuda', "Compute Settings"), { "deep_cache_interval": OptionInfo(3, "DeepCache cache interval", gr.Slider, {"minimum": 1, "maximum": 10, "step": 1}), "quant_sep": OptionInfo("

Model Quantization

", "", gr.HTML), + "quant_shuffle_weights": OptionInfo(False, "Shuffle the weights between GPU and CPU when quantizing"), "diffusers_quantization": OptionInfo(False, "Dynamic quantization with TorchAO"), "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}),