Move post load quant functions to a single function in model_quant

This commit is contained in:
Disty0
2025-04-20 17:04:25 +03:00
parent 707d7483a4
commit 434bb660ce
2 changed files with 14 additions and 8 deletions
+13
View File
@@ -508,3 +508,16 @@ def get_dit_args(load_config:dict={}, module:str=None, device_map:bool=False, al
else:
quant_args = {}
return config, quant_args
def do_post_load_quant(sd_model):
from modules import shared
if shared.opts.nncf_compress_weights and not (shared.opts.cuda_compile and shared.opts.cuda_compile_backend == "openvino_fx"):
sd_model = nncf_compress_weights(sd_model)
if shared.opts.optimum_quanto_weights:
sd_model = optimum_quanto_weights(sd_model)
if shared.opts.torchao_quantization and shared.opts.torchao_quantization_mode == 'post':
sd_model = torchao_quantization(sd_model)
if shared.opts.layerwise_quantization:
apply_layerwise(sd_model)
return sd_model
+1 -8
View File
@@ -146,8 +146,6 @@ def set_diffuser_options(sd_model, vae=None, op:str='model', offload:bool=True,
model.eval()
return model
sd_model = apply_function_to_model(sd_model, eval_model, ["Model", "VAE", "TE"], op="eval")
if len(shared.opts.torchao_quantization) > 0 and shared.opts.torchao_quantization_mode == 'post':
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')
@@ -602,12 +600,7 @@ def load_diffuser(checkpoint_info=None, already_loaded_state_dict=None, timer=No
prompt_parser_diffusers.cache.clear()
set_diffuser_options(sd_model, vae, op, offload=False)
if shared.opts.nncf_compress_weights and not (shared.opts.cuda_compile and shared.opts.cuda_compile_backend == "openvino_fx"):
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 = 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)
sd_model = model_quant.do_post_load_quant(sd_model) # run this before move model so it can be compressed in CPU
timer.record("options")
set_diffuser_offload(sd_model, op)