Public/automatic
1
0
Fork 0
mirror of https://github.com/vladmandic/automatic synced 2026-09-19 09:14:35 +02:00
Code Issues Packages Projects Releases Wiki Activity

Quant add shuffle models option

Browse Source
This commit is contained in:
Disty0
2024-08-04 04:46:06 +03:00
parent 8cc03541cc
commit dc9e60aa67
2 changed files with 71 additions and 24 deletions
Download Patch File Download Diff File Expand all files Collapse all files
modules/sd_models_compile.py
+70 -24
View File
@@ -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),
modules/shared.py
+1
View File
@@ -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("<h2>Model Quantization</h2>", "", 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}),