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

NNCF compression support on CPU and add INT8 option for T5

Browse Source
This commit is contained in:
Disty0
2024-06-19 21:23:47 +03:00
parent 34c7f28aa5
commit bf9565cb46
4 changed files with 70 additions and 28 deletions
Download Patch File Download Diff File Expand all files Collapse all files
modules/model_t5.py
+27 -1
View File
@@ -37,15 +37,41 @@ def load_t5(t5=None, cache_dir=None):
cache_dir=cache_dir,
torch_dtype=devices.dtype,
)
elif 'int8' in t5.lower():
modelloader.hf_login()
from installer import install
install('nncf==2.7.0', quiet=True)
from modules.sd_models_compile import nncf_compress_model
from modules.sd_hijack import NNCF_T5DenseGatedActDense # T5DenseGatedActDense uses fp32
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)):
t5.encoder.block[i].layer[1].DenseReluDense = NNCF_T5DenseGatedActDense(
t5.encoder.block[i].layer[1].DenseReluDense
)
t5 = nncf_compress_model(t5)
else:
t5 = None
return t5
def set_t5(pipe, module, t5=None, cache_dir=None):
from modules import devices
from modules import devices, shared
if pipe is None or not hasattr(pipe, module):
return pipe
t5 = load_t5(t5=t5, cache_dir=cache_dir)
setattr(pipe, module, t5)
if shared.cmd_opts.lowvram or shared.opts.diffusers_seq_cpu_offload:
from accelerate import cpu_offload
getattr(pipe, module).to("cpu")
cpu_offload(getattr(pipe, module), devices.device, offload_buffers=len(getattr(pipe, module)._parameters) > 0) # pylint: disable=protected-access
elif shared.cmd_opts.medvram or shared.opts.diffusers_model_cpu_offload:
if not hasattr(pipe, "_all_hooks") or len(pipe._all_hooks) == 0: # pylint: disable=protected-access
pipe.enable_model_cpu_offload(device=devices.device)
else:
pipe.maybe_free_model_hooks()
devices.torch_gc()
modules/sd_models.py
+21 -10
View File
@@ -653,15 +653,10 @@ def copy_diffuser_options(new_pipe, orig_pipe):
new_pipe.is_sd1 = getattr(orig_pipe, 'is_sd1', True)
def set_diffuser_options(sd_model, vae = None, op: str = 'model'):
def set_diffuser_options(sd_model, vae = None, op: str = 'model', offload=True):
if sd_model is None:
shared.log.warning(f'{op} is not loaded')
return
if (shared.opts.diffusers_model_cpu_offload or shared.cmd_opts.medvram) and (shared.opts.diffusers_seq_cpu_offload or shared.cmd_opts.lowvram):
shared.log.warning(f'Setting {op}: Model CPU offload and Sequential CPU offload are not compatible')
shared.log.debug(f'Setting {op}: disabling model CPU offload')
shared.opts.diffusers_model_cpu_offload=False
shared.cmd_opts.medvram=False
if hasattr(sd_model, "watermark"):
sd_model.watermark = NoWatermark()
@@ -717,6 +712,20 @@ def set_diffuser_options(sd_model, vae = None, op: str = 'model'):
shared.log.debug(f'Setting {op}: enable channels last')
sd_model.unet.to(memory_format=torch.channels_last)
if offload:
set_diffuser_offload(sd_model, op)
def set_diffuser_offload(sd_model, op: str = 'model'):
if sd_model is None:
shared.log.warning(f'{op} is not loaded')
return
if (shared.opts.diffusers_model_cpu_offload or shared.cmd_opts.medvram) and (shared.opts.diffusers_seq_cpu_offload or shared.cmd_opts.lowvram):
shared.log.warning(f'Setting {op}: Model CPU offload and Sequential CPU offload are not compatible')
shared.log.debug(f'Setting {op}: disabling model CPU offload')
shared.opts.diffusers_model_cpu_offload=False
shared.cmd_opts.medvram=False
if not (hasattr(sd_model, "has_accelerate") and sd_model.has_accelerate):
sd_model.has_accelerate = False
if hasattr(sd_model, "enable_model_cpu_offload"):
if shared.cmd_opts.medvram or shared.opts.diffusers_model_cpu_offload:
shared.log.debug(f'Setting {op}: enable model CPU offload')
@@ -1130,7 +1139,12 @@ def load_diffuser(checkpoint_info=None, already_loaded_state_dict=None, timer=No
sd_model.embedding_db.load_textual_inversion_embeddings(force_reload=True)
timer.record("embeddings")
set_diffuser_options(sd_model, vae, op)
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 = sd_models_compile.nncf_compress_weights(sd_model) # run this before move model so it can be compressed in CPU
timer.record("options")
set_diffuser_offload(sd_model, op)
if op == 'model':
sd_vae.apply_vae_config(shared.sd_model.sd_checkpoint_info.filename, vae_file, sd_model)
if op == 'refiner' and shared.opts.diffusers_move_refiner:
@@ -1145,9 +1159,6 @@ def load_diffuser(checkpoint_info=None, already_loaded_state_dict=None, timer=No
if shared.opts.ipex_optimize:
sd_model = sd_models_compile.ipex_optimize(sd_model)
if shared.opts.nncf_compress_weights and not (shared.opts.cuda_compile and shared.opts.cuda_compile_backend == "openvino_fx"):
sd_model = sd_models_compile.nncf_compress_weights(sd_model)
if (shared.opts.cuda_compile and shared.opts.cuda_compile_backend != 'none'):
sd_model = sd_models_compile.compile_diffusers(sd_model)
timer.record("compile")
modules/sd_models_compile.py
+21 -16
View File
@@ -114,27 +114,32 @@ def ipex_optimize(sd_model):
shared.log.warning(f"IPEX Optimize: error: {e}")
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):
import 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)
if hasattr(model, "set_input_embeddings") and backup_embeddings is not None:
model.set_input_embeddings(backup_embeddings)
devices.torch_gc(force=True)
return model
def nncf_compress_weights(sd_model):
try:
t0 = time.time()
if sd_model.device.type == "meta":
shared.log.warning("Compress Weights is not compatible with Sequential CPU offload")
return sd_model
from installer import install
install('nncf==2.7.0', quiet=True)
def nncf_compress_model(model):
return_device = model.device
model.eval()
backup_embeddings = None
if hasattr(model, "get_input_embeddings"):
backup_embeddings = copy.deepcopy(model.get_input_embeddings())
model = nncf.compress_weights(model.to(devices.device)).to(return_device)
if hasattr(model, "set_input_embeddings") and backup_embeddings is not None:
model.set_input_embeddings(backup_embeddings)
devices.torch_gc(force=True)
return model
import nncf
shared.compiled_model_state = CompiledModelState()
shared.compiled_model_state.is_compiled = True
modules/shared.py
+1 -1
View File
@@ -391,7 +391,7 @@ options_templates.update(options_section(('sd', "Execution & Models"), {
"sd_model_refiner": OptionInfo('None', "Refiner model", gr.Dropdown, lambda: {"choices": ['None'] + list_checkpoint_tiles()}, refresh=refresh_checkpoints),
"sd_vae": OptionInfo("Automatic", "VAE model", gr.Dropdown, lambda: {"choices": shared_items.sd_vae_items()}, refresh=shared_items.refresh_vae_list),
"sd_unet": OptionInfo("None", "UNET model", gr.Dropdown, lambda: {"choices": shared_items.sd_unet_items()}, refresh=shared_items.refresh_unet_list),
"sd_text_encoder": OptionInfo('None', "Text encoder model", gr.Dropdown, lambda: {"choices": ['None', 'T5 FP4', 'T5 FP8', 'T5 FP16']}),
"sd_text_encoder": OptionInfo('None', "Text encoder model", gr.Dropdown, lambda: {"choices": ['None', 'T5 FP4', 'T5 FP8', 'T5 INT8', 'T5 FP16']}),
"sd_checkpoint_autoload": OptionInfo(True, "Model autoload on start"),
"sd_model_dict": OptionInfo('None', "Use separate base dict", gr.Dropdown, lambda: {"choices": ['None'] + list_checkpoint_tiles()}, refresh=refresh_checkpoints),
"stream_load": OptionInfo(False, "Load models using stream loading method", gr.Checkbox, {"visible": not native }),