diff --git a/modules/processing_diffusers.py b/modules/processing_diffusers.py index 7e005c358..c28b893bd 100644 --- a/modules/processing_diffusers.py +++ b/modules/processing_diffusers.py @@ -8,7 +8,7 @@ import numpy as np import torch import torchvision.transforms.functional as TF import diffusers -from modules import shared, devices, processing, sd_samplers, sd_models, images, errors, prompt_parser_diffusers, sd_hijack_hypertile, processing_correction, processing_vae +from modules import shared, devices, processing, sd_samplers, sd_models, images, errors, prompt_parser_diffusers, sd_hijack_hypertile, processing_correction, processing_vae, sd_models_compile from modules.processing_helpers import resize_init_images, resize_hires, fix_prompts, calculate_base_steps, calculate_hires_steps, calculate_refiner_steps from modules.onnx_impl import preprocess_pipeline as preprocess_onnx_pipeline, check_parameters_changed as olive_check_parameters_changed @@ -214,7 +214,7 @@ def process_diffusers(p: processing.StableDiffusionProcessing): if hasattr(model, "decoder") and hasattr(model, "prior_prior") and 'prior_num_inference_steps' in possible: steps = kwargs.pop("num_inference_steps", 20) args["prior_num_inference_steps"] = steps - args["num_inference_steps"] = max(int(steps / 2), 1) # TODO: add another slider without overcrowding the UI + args["num_inference_steps"] = max(int(steps / 2), 1) if hasattr(model, "decoder") and hasattr(model, "prior_prior") and 'prior_guidance_scale' in possible: cfg_scale = kwargs.pop("guidance_scale", p.cfg_scale) args["prior_guidance_scale"] = cfg_scale @@ -309,41 +309,6 @@ def process_diffusers(p: processing.StableDiffusionProcessing): debug(f'Diffusers pipeline args: {args}') return args - def recompile_model(hires=False): # recompile if a parameter changes - if shared.opts.cuda_compile and shared.opts.cuda_compile_backend != 'none': - if shared.opts.cuda_compile_backend == "openvino_fx": - compile_height = p.height if not hires and hasattr(p, 'height') else p.hr_upscale_to_y - compile_width = p.width if not hires and hasattr(p, 'width') else p.hr_upscale_to_x - if (shared.compiled_model_state is None or - (not shared.compiled_model_state.first_pass - and (shared.compiled_model_state.height != compile_height - or shared.compiled_model_state.width != compile_width - or shared.compiled_model_state.batch_size != p.batch_size))): - shared.log.info("OpenVINO: Parameter change detected") - shared.log.info("OpenVINO: Recompiling base model") - sd_models.unload_model_weights(op='model') - sd_models.reload_model_weights(op='model') - if is_refiner_enabled(): - shared.log.info("OpenVINO: Recompiling refiner") - sd_models.unload_model_weights(op='refiner') - sd_models.reload_model_weights(op='refiner') - shared.compiled_model_state.height = compile_height - shared.compiled_model_state.width = compile_width - shared.compiled_model_state.batch_size = p.batch_size - - def openvino_post_compile(op="base"): # delete unet after OpenVINO compile - 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 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 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) - def update_sampler(sd_model, second_pass=False): sampler_selection = p.hr_sampler_name if second_pass else p.sampler_name if sd_model.__class__.__name__ in ['AmusedPipeline']: @@ -405,7 +370,7 @@ def process_diffusers(p: processing.StableDiffusionProcessing): sd_models.move_model(shared.sd_model, devices.device) # recompile if a parameter changes - recompile_model() + sd_models_compile.openvino_recompile_model(p, hires=False, refiner=False) # pipeline type is set earlier in processing, but check for sanity is_control = getattr(p, 'is_control', False) is True @@ -444,10 +409,12 @@ def process_diffusers(p: processing.StableDiffusionProcessing): p.extra_generation_params["Sampler Eta"] = shared.opts.scheduler_eta try: t0 = time.time() + sd_models_compile.check_deepcache(enable=True) output = shared.sd_model(**base_args) # pylint: disable=not-callable if isinstance(output, dict): output = SimpleNamespace(**output) - openvino_post_compile(op="base") # only executes on compiled vino models + sd_models_compile.openvino_post_compile(op="base") # only executes on compiled vino models + sd_models_compile.check_deepcache(enable=False) if shared.cmd_opts.profile: t1 = time.time() shared.log.debug(f'Profile: pipeline call: {t1-t0:.2f}') @@ -494,7 +461,7 @@ def process_diffusers(p: processing.StableDiffusionProcessing): output.images = resize_hires(p, latents=output.images) if (latent_scale_mode is not None or p.hr_force) and p.denoising_strength > 0: p.ops.append('hires') - recompile_model(hires=True) + sd_models_compile.openvino_recompile_model(p, hires=True, refiner=False) shared.sd_model = sd_models.set_diffuser_pipe(shared.sd_model, sd_models.DiffusersTaskType.IMAGE_2_IMAGE) if shared.sd_model.__class__.__name__ == "OnnxRawPipeline": shared.sd_model = preprocess_onnx_pipeline(p) @@ -518,10 +485,12 @@ def process_diffusers(p: processing.StableDiffusionProcessing): shared.state.job = 'hires' shared.state.sampling_steps = hires_args['num_inference_steps'] try: + sd_models_compile.check_deepcache(enable=True) output = shared.sd_model(**hires_args) # pylint: disable=not-callable if isinstance(output, dict): output = SimpleNamespace(**output) - openvino_post_compile(op="base") + sd_models_compile.check_deepcache(enable=False) + sd_models_compile.openvino_post_compile(op="base") except AssertionError as e: shared.log.info(e) p.init_images = [] @@ -546,6 +515,7 @@ def process_diffusers(p: processing.StableDiffusionProcessing): sd_models.move_model(shared.sd_refiner, devices.device) p.ops.append('refine') p.is_refiner_pass = True + sd_models_compile.openvino_recompile_model(p, hires=False, refiner=True) shared.sd_model = sd_models.set_diffuser_pipe(shared.sd_model, sd_models.DiffusersTaskType.TEXT_2_IMAGE) shared.sd_refiner = sd_models.set_diffuser_pipe(shared.sd_refiner, sd_models.DiffusersTaskType.IMAGE_2_IMAGE) update_sampler(shared.sd_refiner, second_pass=True) @@ -581,7 +551,7 @@ def process_diffusers(p: processing.StableDiffusionProcessing): refiner_output = shared.sd_refiner(**refiner_args) # pylint: disable=not-callable if isinstance(refiner_output, dict): refiner_output = SimpleNamespace(**refiner_output) - openvino_post_compile(op="refiner") + sd_models_compile.openvino_post_compile(op="refiner") except AssertionError as e: shared.log.info(e) diff --git a/modules/sd_models_compile.py b/modules/sd_models_compile.py index 0fd04cbb8..76bbdd970 100644 --- a/modules/sd_models_compile.py +++ b/modules/sd_models_compile.py @@ -1,7 +1,7 @@ import time import logging import torch -from modules import shared, devices +from modules import shared, devices, sd_models from installer import setup_logging @@ -241,6 +241,14 @@ def compile_torch(sd_model): return sd_model +def check_deepcache(enable: bool): + if deepcache_worker is not None: + if enable: + deepcache_worker.enable() + else: + deepcache_worker.disable() + + def compile_deepcache(sd_model): global deepcache_worker # pylint: disable=global-statement try: @@ -249,11 +257,9 @@ def compile_deepcache(sd_model): shared.log.warning(f'Model compile using deep-cache: {e}') return sd_model t0 = time.time() - if deepcache_worker is not None: - deepcache_worker.disable() + check_deepcache(False) deepcache_worker = DeepCacheSDHelper(pipe=sd_model) deepcache_worker.set_params(cache_interval=shared.opts.deep_cache_interval, cache_branch_id=0) - deepcache_worker.enable() t1 = time.time() shared.log.info(f"Model compile: task=DeepCache config={deepcache_worker.params} time={t1-t0:.2f}") # config={'cache_interval': 3, 'cache_layer_id': 0, 'cache_block_id': 0, 'skip_mode': 'uniform'} time=0.00 @@ -276,8 +282,7 @@ def compile_diffusers(sd_model): elif shared.opts.cuda_compile_backend == 'deep-cache': sd_model = compile_deepcache(sd_model) else: - if deepcache_worker is not None: - deepcache_worker.disable() + check_deepcache(False) sd_model = compile_torch(sd_model) return sd_model @@ -305,3 +310,40 @@ def dynamic_quantization(sd_model): except Exception as e: shared.log.error(f"Model dynamic quantization error: {e}") return sd_model + + +def openvino_recompile_model(p, hires=False, refiner=False): # recompile if a parameter changes + if shared.opts.cuda_compile and shared.opts.cuda_compile_backend != 'none': + if shared.opts.cuda_compile_backend == "openvino_fx": + compile_height = p.height if not hires and hasattr(p, 'height') else p.hr_upscale_to_y + compile_width = p.width if not hires and hasattr(p, 'width') else p.hr_upscale_to_x + if (shared.compiled_model_state is None or + (not shared.compiled_model_state.first_pass + and (shared.compiled_model_state.height != compile_height + or shared.compiled_model_state.width != compile_width + or shared.compiled_model_state.batch_size != p.batch_size))): + if refiner: + shared.log.info("OpenVINO: Recompiling refiner") + sd_models.unload_model_weights(op='refiner') + sd_models.reload_model_weights(op='refiner') + else: + shared.log.info("OpenVINO: Recompiling base model") + sd_models.unload_model_weights(op='model') + sd_models.reload_model_weights(op='model') + shared.compiled_model_state.height = compile_height + shared.compiled_model_state.width = compile_width + shared.compiled_model_state.batch_size = p.batch_size + + +def openvino_post_compile(op="base"): # delete unet after OpenVINO compile + 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 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 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)