diff --git a/modules/intel/openvino/__init__.py b/modules/intel/openvino/__init__.py index 5023a2dcd..c2af24637 100644 --- a/modules/intel/openvino/__init__.py +++ b/modules/intel/openvino/__init__.py @@ -450,7 +450,7 @@ def openvino_fx(subgraph, example_inputs): example_inputs_reordered.append(example_inputs[idx1]) example_inputs = example_inputs_reordered - if dont_use_faketensors: + if dont_use_faketensors or shared.opts.openvino_disable_memory_cleanup: pass else: # Delete unused subgraphs diff --git a/modules/processing_diffusers.py b/modules/processing_diffusers.py index fdf29f082..96183f3f3 100644 --- a/modules/processing_diffusers.py +++ b/modules/processing_diffusers.py @@ -317,12 +317,12 @@ def process_diffusers(p: processing.StableDiffusionProcessing): if shared.opts.cuda_compile and shared.opts.cuda_compile_backend == "openvino_fx": if shared.compiled_model_state.first_pass and op == "base": shared.compiled_model_state.first_pass = False - if hasattr(shared.sd_model, "unet"): + if not shared.opts.openvino_disable_memory_cleanup and hasattr(shared.sd_model, "unet"): shared.sd_model.unet.apply(sd_models.convert_to_faketensors) devices.torch_gc(force=True) if shared.compiled_model_state.first_pass_refiner and op == "refiner": shared.compiled_model_state.first_pass_refiner = False - if hasattr(shared.sd_refiner, "unet"): + if not shared.opts.openvino_disable_memory_cleanup and hasattr(shared.sd_refiner, "unet"): shared.sd_refiner.unet.apply(sd_models.convert_to_faketensors) devices.torch_gc(force=True) diff --git a/modules/processing_vae.py b/modules/processing_vae.py index 55c545cad..985bcdac4 100644 --- a/modules/processing_vae.py +++ b/modules/processing_vae.py @@ -51,7 +51,7 @@ def full_vae_decode(latents, model): # Delete PyTorch VAE after OpenVINO compile if shared.opts.cuda_compile and shared.opts.cuda_compile_backend == "openvino_fx" and shared.compiled_model_state.first_pass_vae: shared.compiled_model_state.first_pass_vae = False - if hasattr(shared.sd_model, "vae"): + if not shared.opts.openvino_disable_memory_cleanup and hasattr(shared.sd_model, "vae"): model.vae.apply(sd_models.convert_to_faketensors) devices.torch_gc(force=True) diff --git a/modules/shared.py b/modules/shared.py index 7082be669..c986333cb 100644 --- a/modules/shared.py +++ b/modules/shared.py @@ -390,6 +390,7 @@ options_templates.update(options_section(('cuda', "Compute Settings"), { "nncf_compress_weights_mode": OptionInfo("INT8", "OpenVINO compress mode for NNCF", gr.Radio, {"choices": ['INT8', 'INT8_SYM', 'INT4_ASYM', 'INT4_SYM', 'NF4'], "visible": cmd_opts.use_openvino}), "nncf_compress_weights_raito": OptionInfo(1.0, "OpenVINO compress ratio for NNCF", gr.Slider, {"minimum": 0, "maximum": 1, "step": 0.01, "visible": cmd_opts.use_openvino}), "openvino_disable_model_caching": OptionInfo(False, "OpenVINO disable model caching", gr.Checkbox, {"visible": cmd_opts.use_openvino}), + "openvino_disable_memory_cleanup": OptionInfo(True, "OpenVINO disable memory cleanup after compile", gr.Checkbox, {"visible": cmd_opts.use_openvino}), "directml_sep": OptionInfo("