diff --git a/modules/model_t5.py b/modules/model_t5.py index 02764ea57..7b735794c 100644 --- a/modules/model_t5.py +++ b/modules/model_t5.py @@ -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() diff --git a/modules/sd_models.py b/modules/sd_models.py index ef18ebf5e..be01df7b3 100644 --- a/modules/sd_models.py +++ b/modules/sd_models.py @@ -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") diff --git a/modules/sd_models_compile.py b/modules/sd_models_compile.py index a4b4fb8af..6e629788c 100644 --- a/modules/sd_models_compile.py +++ b/modules/sd_models_compile.py @@ -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 diff --git a/modules/shared.py b/modules/shared.py index d65084053..fcdf9bbb8 100644 --- a/modules/shared.py +++ b/modules/shared.py @@ -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 }),