mirror of
https://github.com/vladmandic/automatic
synced 2026-09-19 09:14:35 +02:00
Move quant functions to model_quant.py
This commit is contained in:
Disty0
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user