From 9b579bfd96d82e5e9cfc9090ff14a67ee36ca4ef Mon Sep 17 00:00:00 2001 From: Disty0 Date: Thu, 23 Jan 2025 21:50:26 +0300 Subject: [PATCH] Move quant functions to model_quant.py --- modules/model_quant.py | 237 ++++++++++++++++++++++++++++-- modules/model_te.py | 4 +- modules/sd_models.py | 10 +- modules/sd_models_compile.py | 276 +---------------------------------- modules/sd_models_utils.py | 83 ++++++++++- 5 files changed, 318 insertions(+), 292 deletions(-) diff --git a/modules/model_quant.py b/modules/model_quant.py index 0a56805df..24c517dc3 100644 --- a/modules/model_quant.py +++ b/modules/model_quant.py @@ -1,11 +1,17 @@ import sys +import copy +import time import diffusers -from installer import install, log +from installer import install, log, setup_logging -bnb = None -quanto = None ao = None +bnb = None +intel_nncf = None +optimum_quanto = None + +quant_last_model_name = None +quant_last_model_device = None def get_quant(name): @@ -114,22 +120,42 @@ def load_bnb(msg='', silent=False): def load_quanto(msg='', silent=False): from modules import shared - global quanto # pylint: disable=global-statement - if quanto is not None: - return quanto + global optimum_quanto # pylint: disable=global-statement + if optimum_quanto is not None: + return optimum_quanto install('optimum-quanto==0.2.6', quiet=True) try: - from optimum import quanto as optimum_quanto # pylint: disable=no-name-in-module - quanto = optimum_quanto + 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 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') - return quanto + return optimum_quanto except Exception as e: if len(msg) > 0: log.error(f"{msg} failed to import optimum.quanto: {e}") - quanto = None + optimum_quanto = None + if not silent: + raise + return None + + +def load_nncf(msg='', silent=False): + global intel_nncf # pylint: disable=global-statement + if intel_nncf is not None: + return intel_nncf + install('nncf==2.7.0', quiet=True) + try: + import nncf + intel_nncf = nncf + 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=nncf version={nncf.__version__} fn={fn}') # pylint: disable=protected-access + return intel_nncf + except Exception as e: + if len(msg) > 0: + log.error(f"{msg} failed to import nncf: {e}") + intel_nncf = None if not silent: raise return None @@ -175,3 +201,194 @@ def apply_layerwise(sd_model, quiet:bool=False): 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}') + + +def nncf_send_to_device(model, device): + for child in model.children(): + if child.__class__.__name__ == "WeightsDecompressor": + child.scale = child.scale.to(device) + child.zero_point = child.zero_point.to(device) + nncf_send_to_device(child, device) + + +def nncf_compress_model(model, op=None, sd_model=None): + from modules import devices, shared + global quant_last_model_name, quant_last_model_device # pylint: disable=global-statement + nncf = load_nncf('Quantize model: type=NNCF') + model.eval() + backup_embeddings = None + if hasattr(model, "get_input_embeddings"): + backup_embeddings = copy.deepcopy(model.get_input_embeddings()) + model = nncf.compress_weights(model) + nncf_send_to_device(model, devices.device) + 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) + 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) + return model + + +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}") + 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") + 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"Quantization: type=NNCF time={t1-t0:.2f}") + except Exception as e: + shared.log.warning(f"Quantization: type=NNCF {e}") + 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 sd_model is not None and "Flux" in sd_model.__class__.__name__: # 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"] + 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.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) + 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) + 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"}: + shared.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}") + 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) + + 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) + 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) + + t1 = time.time() + shared.log.info(f"Quantization: type=Optimum.quanto time={t1-t0:.2f}") + except Exception as e: + shared.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: + shared.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}") + 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}") + except Exception as e: + shared.log.error(f"Quantization: type=TorchAO {e}") + setup_logging() # torchao uses dynamo which messes with logging so reset is needed + return sd_model diff --git a/modules/model_te.py b/modules/model_te.py index 16bb6d222..024dda47c 100644 --- a/modules/model_te.py +++ b/modules/model_te.py @@ -64,12 +64,12 @@ def load_t5(name=None, cache_dir=None): t5 = transformers.T5EncoderModel.from_pretrained(repo_id, subfolder='text_encoder_3', quantization_config=quantization_config, cache_dir=cache_dir, torch_dtype=devices.dtype) elif 'qint8' in name.lower(): model_quant.load_quanto('Load model: type=T5') - from modules.sd_models_compile import optimum_quanto_model + from modules.model_quant import optimum_quanto_model t5 = transformers.T5EncoderModel.from_pretrained(repo_id, subfolder='text_encoder_3', cache_dir=cache_dir, torch_dtype=devices.dtype) t5 = optimum_quanto_model(t5, weights="qint8", activations="none") elif 'int8' in name.lower(): install('nncf==2.7.0', quiet=True) - from modules.sd_models_compile import nncf_compress_model + from modules.model_quant import nncf_compress_model from modules.sd_hijack import NNCF_T5DenseGatedActDense t5 = transformers.T5EncoderModel.from_pretrained(repo_id, subfolder='text_encoder_3', cache_dir=cache_dir, torch_dtype=devices.dtype) for i in range(len(t5.encoder.block)): diff --git a/modules/sd_models.py b/modules/sd_models.py index 0b98bb504..df49b4918 100644 --- a/modules/sd_models.py +++ b/modules/sd_models.py @@ -16,7 +16,7 @@ from modules.modeldata import model_data from modules.sd_checkpoint import CheckpointInfo, select_checkpoint, list_models, checkpoints_list, checkpoint_titles, get_closet_checkpoint_match, model_hash, update_model_hashes, setup_model, write_metadata, read_metadata_from_safetensors # pylint: disable=unused-import from modules.sd_offload import disable_offload, set_diffuser_offload, apply_balanced_offload, set_accelerate # pylint: disable=unused-import from modules.sd_models_legacy import get_checkpoint_state_dict, load_model_weights, load_model, repair_config # pylint: disable=unused-import -from modules.sd_models_utils import NoWatermark, get_signature, get_call, path_to_repo, patch_diffuser_config, convert_to_faketensors, read_state_dict, get_state_dict_from_checkpoint # pylint: disable=unused-import +from modules.sd_models_utils import NoWatermark, get_signature, get_call, path_to_repo, patch_diffuser_config, convert_to_faketensors, read_state_dict, get_state_dict_from_checkpoint, apply_function_to_model # pylint: disable=unused-import model_dir = "Stable-diffusion" @@ -130,9 +130,9 @@ def set_diffuser_options(sd_model, vae=None, op:str='model', offload:bool=True, model.requires_grad_(False) model.eval() return model - sd_model = sd_models_compile.apply_compile_to_model(sd_model, eval_model, ["Model", "VAE", "Text Encoder"], op="eval") + sd_model = apply_function_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': - sd_model = sd_models_compile.torchao_quantization(sd_model) + sd_model = model_quant.torchao_quantization(sd_model) if shared.opts.opt_channelslast and hasattr(sd_model, 'unet'): shared.log.quiet(quiet, f'Setting {op}: channels-last=True') @@ -567,9 +567,9 @@ def load_diffuser(checkpoint_info=None, already_loaded_state_dict=None, timer=No set_diffuser_options(sd_model, vae, op, offload=False) if shared.opts.nncf_compress_weights and not ('Model' in shared.opts.cuda_compile and shared.opts.cuda_compile_backend == "openvino_fx"): - sd_model = sd_models_compile.nncf_compress_weights(sd_model) # run this before move model so it can be compressed in CPU + sd_model = model_quant.nncf_compress_weights(sd_model) # run this before move model so it can be compressed in CPU if shared.opts.optimum_quanto_weights: - sd_model = sd_models_compile.optimum_quanto_weights(sd_model) # run this before move model so it can be compressed in CPU + sd_model = model_quant.optimum_quanto_weights(sd_model) # run this before move model so it can be compressed in CPU if shared.opts.layerwise_quantization: model_quant.apply_layerwise(sd_model) timer.record("options") diff --git a/modules/sd_models_compile.py b/modules/sd_models_compile.py index 0972c3350..6754984e5 100644 --- a/modules/sd_models_compile.py +++ b/modules/sd_models_compile.py @@ -25,90 +25,9 @@ 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, op="unet", sd_model=sd_model) - if hasattr(sd_model, 'transformer') and hasattr(sd_model.transformer, 'config'): - 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, 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, 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 "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') and hasattr(sd_model.decoder_pipe.text_encoder, 'config'): - 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 - for i in range(len(sd_model.text_encoder.encoder.block)): - sd_model.text_encoder.encoder.block[i].layer[1].DenseReluDense = NNCF_T5DenseGatedActDense( - 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, 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 - for i in range(len(sd_model.text_encoder_2.encoder.block)): - sd_model.text_encoder_2.encoder.block[i].layer[1].DenseReluDense = NNCF_T5DenseGatedActDense( - 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, 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 - for i in range(len(sd_model.text_encoder_3.encoder.block)): - sd_model.text_encoder_3.encoder.block[i].layer[1].DenseReluDense = NNCF_T5DenseGatedActDense( - 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, op="text_encoder_3", sd_model=sd_model) - if hasattr(sd_model, 'prior_pipe') and hasattr(sd_model.prior_pipe, 'text_encoder') and hasattr(sd_model.prior_pipe.text_encoder, 'config'): - 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'): - if op == "compile": - sd_model.vae.decode = function(sd_model.vae.decode, op="vae_decode", sd_model=sd_model) - sd_model.vae.encode = function(sd_model.vae.encode, op="vae_encode", sd_model=sd_model) - else: - sd_model.vae = function(sd_model.vae, op="vae", sd_model=sd_model) - if hasattr(sd_model, 'movq') and hasattr(sd_model.movq, 'decode'): - if op == "compile": - sd_model.movq.decode = function(sd_model.movq.decode, op="movq_decode", sd_model=sd_model) - sd_model.movq.encode = function(sd_model.movq.encode, op="movq_encode", sd_model=sd_model) - else: - sd_model.movq = function(sd_model.movq, op="movq", sd_model=sd_model) - if hasattr(sd_model, 'vqgan') and hasattr(sd_model.vqgan, 'decode'): - if op == "compile": - sd_model.vqgan.decode = function(sd_model.vqgan.decode, op="vqgan_decode", sd_model=sd_model) - sd_model.vqgan.encode = function(sd_model.vqgan.encode, op="vqgan_encode", sd_model=sd_model) - else: - 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'): - if op == "compile": - sd_model.decoder_pipe.vqgan.decode = function(sd_model.decoder_pipe.vqgan.decode, op="vqgan_decode", sd_model=sd_model) - sd_model.decoder_pipe.vqgan.encode = function(sd_model.decoder_pipe.vqgan.encode, op="vqgan_encode", sd_model=sd_model) - else: - 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 - - def ipex_optimize(sd_model): try: t0 = time.time() @@ -133,7 +52,7 @@ def ipex_optimize(sd_model): devices.torch_gc() return model - sd_model = apply_compile_to_model(sd_model, ipex_optimize_model, shared.opts.ipex_optimize, op="ipex") + sd_model = sd_models.apply_function_to_model(sd_model, ipex_optimize_model, shared.opts.ipex_optimize, op="ipex") t1 = time.time() shared.log.info(f"IPEX Optimize: time={t1-t0:.2f}") @@ -142,169 +61,6 @@ def ipex_optimize(sd_model): return sd_model -def nncf_send_to_device(model): - for child in model.children(): - if child.__class__.__name__ == "WeightsDecompressor": - child.scale = child.scale.to(devices.device) - child.zero_point = child.zero_point.to(devices.device) - nncf_send_to_device(child) - - -def nncf_compress_model(model, op=None, sd_model=None): - import nncf - global quant_last_model_name, quant_last_model_device # pylint: disable=global-statement - model.eval() - backup_embeddings = None - if hasattr(model, "get_input_embeddings"): - backup_embeddings = copy.deepcopy(model.get_input_embeddings()) - model = nncf.compress_weights(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) - 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) - return model - - -def nncf_compress_weights(sd_model): - try: - t0 = time.time() - shared.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 - 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"Quantization: type=NNCF time={t1-t0:.2f}") - except Exception as e: - shared.log.warning(f"Quantization: type=NNCF {e}") - return sd_model - - -def optimum_quanto_model(model, op=None, sd_model=None, weights=None, activations=None): - quanto = model_quant.load_quanto('Compile model: type=Optimum Quanto') - global quant_last_model_name, quant_last_model_device # pylint: disable=global-statement - if sd_model is not None and "Flux" in sd_model.__class__.__name__: # 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"] - 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.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) - 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) - return model - - -def optimum_quanto_weights(sd_model): - try: - 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") - return sd_model - t0 = time.time() - shared.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 = model_quant.load_quanto() - quanto.tensor.qbits.QBitsTensor.create = lambda *args, **kwargs: quanto.tensor.qbits.QBitsTensor(*args, **kwargs) - - 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 - - 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 = 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) - if hasattr(sd_model, "encode_prompt"): - sd_model.encode_prompt = original_encode_prompt - devices.torch_gc(force=True) - - t1 = time.time() - shared.log.info(f"Quantization: type=Optimum.quanto time={t1-t0:.2f}") - except Exception as e: - shared.log.warning(f"Quantization: type=Optimum.quanto {e}") - return sd_model - - def optimize_openvino(sd_model): try: from modules.intel.openvino import openvino_fx # pylint: disable=unused-import @@ -444,7 +200,7 @@ def compile_torch(sd_model): except Exception as e: shared.log.error(f"Model compile: torch inductor config error: {e}") - sd_model = apply_compile_to_model(sd_model, function=torch_compile_model, options=shared.opts.cuda_compile, op="compile") + sd_model = sd_models.apply_function_to_model(sd_model, function=torch_compile_model, options=shared.opts.cuda_compile, op="compile") setup_logging() # compile messes with logging so reset is needed if shared.opts.cuda_compile_precompile: @@ -503,34 +259,6 @@ def compile_diffusers(sd_model): return sd_model -def torchao_quantization(sd_model): - try: - 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 - - 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 - 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}") - try: - t0 = time.time() - apply_compile_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}") - except Exception as e: - shared.log.error(f"Quantization: type=TorchAO {e}") - setup_logging() # torchao uses dynamo which messes with logging so reset is needed - return sd_model - - def openvino_recompile_model(p, hires=False, refiner=False): # recompile if a parameter changes # pylint: disable=unused-argument if shared.opts.cuda_compile_backend == "openvino_fx" and 'Model' in shared.opts.cuda_compile: compile_height = p.height if not hires and hasattr(p, 'height') else p.hr_upscale_to_y diff --git a/modules/sd_models_utils.py b/modules/sd_models_utils.py index 0ff903483..a09ca73a7 100644 --- a/modules/sd_models_utils.py +++ b/modules/sd_models_utils.py @@ -1,4 +1,5 @@ import io +import copy import json import inspect import os.path @@ -6,7 +7,7 @@ from rich import progress # pylint: disable=redefined-builtin import torch import safetensors.torch -from modules import paths, shared, errors +from modules import paths, shared, devices, errors from modules.sd_checkpoint import CheckpointInfo, select_checkpoint, list_models, checkpoints_list, checkpoint_titles, get_closet_checkpoint_match, model_hash, update_model_hashes, setup_model, write_metadata, read_metadata_from_safetensors # pylint: disable=unused-import from modules.sd_offload import disable_offload, set_diffuser_offload, apply_balanced_offload, set_accelerate # pylint: disable=unused-import from modules.sd_models_legacy import get_checkpoint_state_dict, load_model_weights, load_model, repair_config # pylint: disable=unused-import @@ -149,3 +150,83 @@ def patch_diffuser_config(sd_model, model_file): component.config[k] = v updated[k] = v return sd_model + + +def apply_function_to_model(sd_model, function, options, op=None): + if "Model" in options or "Transformer" in options: + if hasattr(sd_model, 'transformer') and hasattr(sd_model.transformer, 'config'): + sd_model.transformer = function(sd_model.transformer, op="transformer", sd_model=sd_model) + if "Model" in options: + if hasattr(sd_model, 'unet') and hasattr(sd_model.unet, 'config'): + sd_model.unet = function(sd_model.unet, op="unet", 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, 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, 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 "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') and hasattr(sd_model.decoder_pipe.text_encoder, 'config'): + 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 + for i in range(len(sd_model.text_encoder.encoder.block)): + sd_model.text_encoder.encoder.block[i].layer[1].DenseReluDense = NNCF_T5DenseGatedActDense( + 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, 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 + for i in range(len(sd_model.text_encoder_2.encoder.block)): + sd_model.text_encoder_2.encoder.block[i].layer[1].DenseReluDense = NNCF_T5DenseGatedActDense( + 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, 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 + for i in range(len(sd_model.text_encoder_3.encoder.block)): + sd_model.text_encoder_3.encoder.block[i].layer[1].DenseReluDense = NNCF_T5DenseGatedActDense( + 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, op="text_encoder_3", sd_model=sd_model) + if hasattr(sd_model, 'prior_pipe') and hasattr(sd_model.prior_pipe, 'text_encoder') and hasattr(sd_model.prior_pipe.text_encoder, 'config'): + 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'): + if op == "compile": + sd_model.vae.decode = function(sd_model.vae.decode, op="vae_decode", sd_model=sd_model) + sd_model.vae.encode = function(sd_model.vae.encode, op="vae_encode", sd_model=sd_model) + else: + sd_model.vae = function(sd_model.vae, op="vae", sd_model=sd_model) + if hasattr(sd_model, 'movq') and hasattr(sd_model.movq, 'decode'): + if op == "compile": + sd_model.movq.decode = function(sd_model.movq.decode, op="movq_decode", sd_model=sd_model) + sd_model.movq.encode = function(sd_model.movq.encode, op="movq_encode", sd_model=sd_model) + else: + sd_model.movq = function(sd_model.movq, op="movq", sd_model=sd_model) + if hasattr(sd_model, 'vqgan') and hasattr(sd_model.vqgan, 'decode'): + if op == "compile": + sd_model.vqgan.decode = function(sd_model.vqgan.decode, op="vqgan_decode", sd_model=sd_model) + sd_model.vqgan.encode = function(sd_model.vqgan.encode, op="vqgan_encode", sd_model=sd_model) + else: + 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'): + if op == "compile": + sd_model.decoder_pipe.vqgan.decode = function(sd_model.decoder_pipe.vqgan.decode, op="vqgan_decode", sd_model=sd_model) + sd_model.decoder_pipe.vqgan.encode = function(sd_model.decoder_pipe.vqgan.encode, op="vqgan_encode", sd_model=sd_model) + else: + 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