diff --git a/extensions-builtin/Lora/lora.py b/extensions-builtin/Lora/lora.py index 02bb2092c..f23bc6b76 100644 --- a/extensions-builtin/Lora/lora.py +++ b/extensions-builtin/Lora/lora.py @@ -239,6 +239,22 @@ def load_loras(names, multipliers=None): failed_to_load_loras = [] + recompile_model = False + if shared.opts.cuda_compile and shared.opts.cuda_compile_backend == "openvino_fx": + if len(names) == len(shared.compiled_model_state.lora_model): + for i, name in enumerate(names): + if shared.compiled_model_state.lora_model[i] != f"{name}:{multipliers[i]}": + recompile_model = True + break + else: + recompile_model = True + shared.compiled_model_state.lora_model = [] + if recompile_model: + sd_models.unload_model_weights(op='model') + shared.opts.cuda_compile = False + sd_models.reload_model_weights(op='model') + shared.opts.cuda_compile = True + for i, name in enumerate(names): lora = already_loaded.get(name, None) if shared.backend == shared.Backend.ORIGINAL else None lora_on_disk = loras_on_disk[i] @@ -265,6 +281,10 @@ def load_loras(names, multipliers=None): if len(failed_to_load_loras) > 0: sd_hijack.model_hijack.comments.append("Failed to find Loras: " + ", ".join(failed_to_load_loras)) + + if recompile_model: + shared.log.info("Lora: Recompiling model") + shared.sd_model = sd_models.compile_diffusers(shared.sd_model) def lora_calc_updown(lora, module, target): diff --git a/modules/intel/openvino/__init__.py b/modules/intel/openvino/__init__.py index 6162b674b..88b43610d 100644 --- a/modules/intel/openvino/__init__.py +++ b/modules/intel/openvino/__init__.py @@ -1,23 +1,50 @@ import os import torch -from openvino.frontend.pytorch.torchdynamo.execute import execute, partitioned_modules, compiled_cache +from openvino.frontend import FrontEndManager +from openvino.frontend.pytorch.fx_decoder import TorchFXPythonDecoder from openvino.frontend.pytorch.torchdynamo.partition import Partitioner -from openvino.runtime import Core, Type, PartialShape +from openvino.runtime import Core, Type, PartialShape, serialize from torch._dynamo.backends.common import fake_tensor_unsupported from torch._dynamo.backends.registry import register_backend from torch.fx.experimental.proxy_tensor import make_fx -from torch._inductor.compile_fx import compile_fx from torch.utils._pytree import tree_flatten +from types import MappingProxyType from hashlib import sha256 import functools from modules import shared, devices -def openvino_clear_caches(): - global partitioned_modules - global compiled_cache +compiled_cache = {} +max_openvino_partitions = 0 +partitioned_modules = {} - compiled_cache.clear() - partitioned_modules.clear() +DEFAULT_OPENVINO_PYTHON_CONFIG = MappingProxyType( + { + "use_python_fusion_cache": True, + "allow_single_op_fusion": True, + }, +) + +class OpenVINOGraphModule(torch.nn.Module): + def __init__(self, gm, partition_id, use_python_fusion_cache, model_hash_str: str = None, file_name=""): + super().__init__() + self.gm = gm + self.partition_id = partition_id + self.executor_parameters = {"use_python_fusion_cache": use_python_fusion_cache, + "model_hash_str": model_hash_str} + self.file_name = file_name + self.perm_fallback = False + + def __call__(self, *args): + #if self.perm_fallback: + # return self.gm(*args) + + #try: + result = openvino_execute(self.gm, *args, executor_parameters=self.executor_parameters, partition_id=self.partition_id, file_name=self.file_name) + #except Exception: + # self.perm_fallback = True + # return self.gm(*args) + + return result def get_device(): core = Core() @@ -74,6 +101,107 @@ def cached_model_name(model_hash_str, device, args, cache_root, reversed = False return file_name +def check_fully_supported(self, graph_module): + num_fused = 0 + for node in graph_module.graph.nodes: + if node.op == "call_module" and "fused_" in node.name: + num_fused += 1 + elif node.op != "placeholder" and node.op != "output": + return False + if num_fused == 1: + return True + return False + +Partitioner.check_fully_supported = functools.partial(check_fully_supported, Partitioner) + +def execute( + gm, + *args, + executor = "openvino", + executor_parameters = None, + file_name = "" +): + if executor == "openvino": + return openvino_execute_partitioned(gm, *args, executor_parameters=executor_parameters, file_name=file_name) + elif executor == "strictly_openvino": + return openvino_execute(gm, *args, executor_parameters=executor_parameters, file_name=file_name) + + msg = "Received unexpected value for 'executor': {0}. Allowed values are: openvino, strictly_openvino.".format(executor) + raise ValueError(msg) + +def execute_cached(compiled_model, *args): + flat_args, _ = tree_flatten(args) + ov_inputs = [a.detach().cpu().numpy() for a in flat_args] + + if (shared.compiled_model_state.cn_model == []): + ov_inputs.reverse() + + res = compiled_model(ov_inputs) + result = [torch.from_numpy(res[out]) for out in compiled_model.outputs] + return result + +def openvino_clear_caches(): + global partitioned_modules + global compiled_cache + + compiled_cache.clear() + partitioned_modules.clear() + +def openvino_compile(gm, *args, model_hash_str: str = None, file_name=""): + core = Core() + + device = get_device() + cache_root = cache_root_path() + + if file_name is not None and os.path.isfile(file_name + ".xml") and os.path.isfile(file_name + ".bin"): + om = core.read_model(file_name + ".xml") + else: + fe_manager = FrontEndManager() + fe = fe_manager.load_by_framework("pytorch") + + input_shapes = [] + input_types = [] + for input_data in args: + input_types.append(input_data.type()) + input_shapes.append(input_data.size()) + + decoder = TorchFXPythonDecoder(gm, gm, input_shapes=input_shapes, input_types=input_types) + + im = fe.load(decoder) + + om = fe.convert(im) + + if (file_name is not None): + serialize(om, file_name + ".xml", file_name + ".bin") + if (shared.compiled_model_state.cn_model != []): + f = open(file_name + ".txt", "w") + for input_data in args: + f.write(str(input_data.size())) + f.write("\n") + f.close() + + dtype_mapping = { + torch.float32: Type.f32, + torch.float64: Type.f64, + torch.float16: Type.f16, + torch.int64: Type.i64, + torch.int32: Type.i32, + torch.uint8: Type.u8, + torch.int8: Type.i8, + torch.bool: Type.boolean + } + + for idx, input_data in enumerate(args): + om.inputs[idx].get_node().set_element_type(dtype_mapping[input_data.dtype]) + om.inputs[idx].get_node().set_partial_shape(PartialShape(list(input_data.shape))) + om.validate_nodes_and_infer_types() + + if model_hash_str is not None: + core.set_property({'CACHE_DIR': cache_root + '/blob'}) + + compiled = core.compile_model(om, device) + return compiled + def openvino_compile_cached_model(cached_model_path, *example_inputs): core = Core() om = core.read_model(cached_model_path + ".xml") @@ -100,46 +228,92 @@ def openvino_compile_cached_model(cached_model_path, *example_inputs): return compiled_model -def execute_cached(compiled_model, *args): - model_state = shared.compiled_model_state +def openvino_execute(gm, *args, executor_parameters=None, partition_id, file_name=""): + executor_parameters = executor_parameters or DEFAULT_OPENVINO_PYTHON_CONFIG + + use_cache = executor_parameters.get( + "use_python_fusion_cache", + DEFAULT_OPENVINO_PYTHON_CONFIG["use_python_fusion_cache"], + ) + global compiled_cache + + model_hash_str = executor_parameters.get("model_hash_str", None) + if model_hash_str is not None: + model_hash_str = model_hash_str + str(partition_id) + + if use_cache and (partition_id in compiled_cache): + compiled = compiled_cache[partition_id] + else: + if (shared.compiled_model_state.cn_model != [] and file_name is not None + and os.path.isfile(file_name + ".xml") and os.path.isfile(file_name + ".bin")): + compiled = openvino_compile_cached_model(file_name, *args) + else: + compiled = openvino_compile(gm, *args, model_hash_str=model_hash_str, file_name=file_name) + compiled_cache[partition_id] = compiled + flat_args, _ = tree_flatten(args) ov_inputs = [a.detach().cpu().numpy() for a in flat_args] - if (model_state.cn_model == "None"): - ov_inputs.reverse() + res = compiled(ov_inputs) - res = compiled_model(ov_inputs) - result = [torch.from_numpy(res[out]) for out in compiled_model.outputs] - return result + results1 = [torch.from_numpy(res[out]) for out in compiled.outputs] + if len(results1) == 1: + return results1[0] + return results1 -def check_fully_supported(self, graph_module): - num_fused = 0 - for node in graph_module.graph.nodes: +def openvino_execute_partitioned(gm, *args, executor_parameters=None, file_name=""): + executor_parameters = executor_parameters or DEFAULT_OPENVINO_PYTHON_CONFIG + + global partitioned_modules + + use_python_fusion_cache = executor_parameters.get( + "use_python_fusion_cache", + DEFAULT_OPENVINO_PYTHON_CONFIG["use_python_fusion_cache"], + ) + model_hash_str = executor_parameters.get("model_hash_str", None) + + signature = str(id(gm)) + for idx, input_data in enumerate(args): + if isinstance(input_data, torch.Tensor): + signature = signature + "_" + str(idx) + ":" + str(input_data.type())[6:] + ":" + str(input_data.size())[11:-1].replace(" ", "") + else: + signature = signature + "_" + str(idx) + ":" + type(input_data).__name__ + ":val(" + str(input_data) + ")" + + if signature not in partitioned_modules: + partitioned_modules[signature] = partition_graph(gm, use_python_fusion_cache=use_python_fusion_cache, + model_hash_str=model_hash_str, file_name=file_name) + + return partitioned_modules[signature](*args) + +def partition_graph(gm, use_python_fusion_cache: bool, model_hash_str: str = None, file_name=""): + global max_openvino_partitions + for node in gm.graph.nodes: if node.op == "call_module" and "fused_" in node.name: - num_fused += 1 - elif node.op != "placeholder" and node.op != "output": - return False - if num_fused == 1: - return True - return False + openvino_submodule = getattr(gm, node.name) + gm.delete_submodule(node.target) + gm.add_submodule( + node.target, + OpenVINOGraphModule(openvino_submodule, shared.compiled_model_state.partition_id, use_python_fusion_cache, + model_hash_str=model_hash_str, file_name=file_name), + ) + shared.compiled_model_state.partition_id = shared.compiled_model_state.partition_id + 1 -Partitioner.check_fully_supported = functools.partial(check_fully_supported, Partitioner) + return gm @register_backend @fake_tensor_unsupported def openvino_fx(subgraph, example_inputs): - model_state = shared.compiled_model_state executor_parameters = None inputs_reversed = False if os.getenv("OPENVINO_TORCH_MODEL_CACHING") != "0": os.environ.setdefault('OPENVINO_TORCH_MODEL_CACHING', "1") # Create a hash to be used for caching model_hash_str = sha256(subgraph.code.encode('utf-8')).hexdigest() - if (model_state.cn_model != "None" and model_state.partition_id == 0): - model_hash_str = model_hash_str + model_state.cn_model + if (shared.compiled_model_state.cn_model != [] and shared.compiled_model_state.partition_id == 0): + model_hash_str = model_hash_str + str(shared.compiled_model_state.cn_model) - if (model_state.lora_model != "None"): - model_hash_str = model_hash_str + model_state.lora_model + if (shared.compiled_model_state.lora_model != []): + model_hash_str = model_hash_str + str(shared.compiled_model_state.lora_model) executor_parameters = {"model_hash_str": model_hash_str} # Check if the model was fully supported and already cached @@ -148,7 +322,7 @@ def openvino_fx(subgraph, example_inputs): maybe_fs_cached_name = cached_model_name(model_hash_str + "_fs", get_device(), example_inputs, cache_root_path()) if os.path.isfile(maybe_fs_cached_name + ".xml") and os.path.isfile(maybe_fs_cached_name + ".bin"): - if (model_state.cn_model != "None" and model_state.cn_model in maybe_fs_cached_name): + if (shared.compiled_model_state.cn_model != [] and str(shared.compiled_model_state.cn_model) in maybe_fs_cached_name): example_inputs_reordered = [] if (os.path.isfile(maybe_fs_cached_name + ".txt")): f = open(maybe_fs_cached_name + ".txt", "r") @@ -164,7 +338,7 @@ def openvino_fx(subgraph, example_inputs): compiled_model = openvino_compile_cached_model(maybe_fs_cached_name, *example_inputs) def _call(*args): - if (model_state.cn_model != "None" and model_state.cn_model in maybe_fs_cached_name): + if (shared.compiled_model_state.cn_model != [] and str(shared.compiled_model_state.cn_model) in maybe_fs_cached_name): args_reordered = [] if (os.path.isfile(maybe_fs_cached_name + ".txt")): f = open(maybe_fs_cached_name + ".txt", "r") @@ -177,11 +351,11 @@ def openvino_fx(subgraph, example_inputs): args = args_reordered res = execute_cached(compiled_model, *args) - model_state.partition_id = model_state.partition_id + 1 + shared.compiled_model_state.partition_id = shared.compiled_model_state.partition_id + 1 return res return _call else: - maybe_fs_cached_name = None + maybe_fs_cached_name = "" if inputs_reversed: example_inputs.reverse() @@ -202,6 +376,6 @@ def openvino_fx(subgraph, example_inputs): def _call(*args): res = execute(compiled_model, *args, executor="openvino", - executor_parameters=executor_parameters) #, file_name=maybe_fs_cached_name) + executor_parameters=executor_parameters, file_name=maybe_fs_cached_name) return res return _call diff --git a/modules/lora_diffusers.py b/modules/lora_diffusers.py index 72af05722..c6ce96d1e 100644 --- a/modules/lora_diffusers.py +++ b/modules/lora_diffusers.py @@ -35,8 +35,6 @@ def unload_diffusers_lora(): lora_state['loaded'].clear() lora_state['all_loras'] = [] lora_state['multiplier'] = [] - if shared.opts.cuda_compile and shared.opts.cuda_compile_backend == "openvino_fx": - shared.compiled_model_state.lora_model = "None" except Exception as e: shared.log.error(f"LoRA unload failed: {e}") @@ -58,6 +56,8 @@ def load_diffusers_lora(name, lora, strength = 1.0, num_loras = 1): pipe.fuse_lora(lora_scale=strength) fuse = time.time() - t2 lora_state['loaded'].append(f'{lora.filename}:{strength}') + if shared.compiled_model_state is not None: #filename breaks caching + shared.compiled_model_state.lora_model.append(f'{name}:{strength}') else: from safetensors.torch import load_file lora_sd = load_file(lora.filename) @@ -74,10 +74,10 @@ def load_diffusers_lora(name, lora, strength = 1.0, num_loras = 1): lora_network.apply_to(multiplier=strength) lora_state['all_loras'].append(lora_network) lora_state['loaded'].append(f'{lora.filename}:{strength}') + if shared.compiled_model_state is not None: #filename breaks caching + shared.compiled_model_state.lora_model.append(f'{name}:{strength}') t1 = time.time() fuse = f'fuse={fuse:.2f}s' if fuse > 0 else '' - if shared.opts.cuda_compile and shared.opts.cuda_compile_backend == "openvino_fx": - shared.compiled_model_state.lora_model = str(lora_state['loaded']) shared.log.info(f'LoRA loaded: {name} strength={strength} loader="{shared.opts.diffusers_lora_loader}" lora={t1-t0:.2f}s {fuse}') except Exception as e: lines = str(e).splitlines() diff --git a/modules/processing_diffusers.py b/modules/processing_diffusers.py index 6b03640f3..c52193e5b 100644 --- a/modules/processing_diffusers.py +++ b/modules/processing_diffusers.py @@ -260,10 +260,12 @@ def process_diffusers(p: StableDiffusionProcessing, seeds, prompts, negative_pro if shared.opts.cuda_compile_backend == "openvino_fx": compile_height = p.height if not hires else p.hr_upscale_to_y compile_width = p.width if not hires 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 + 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: Resolution change detected") + 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') diff --git a/modules/sd_models.py b/modules/sd_models.py index f27068914..36f00d0af 100644 --- a/modules/sd_models.py +++ b/modules/sd_models.py @@ -123,8 +123,8 @@ class CompiledModelState: self.width = 512 self.batch_size = 1 self.partition_id = 0 - self.cn_model = "None" - self.lora_model = "None" + self.cn_model = [] + self.lora_model = [] class NoWatermark: @@ -657,6 +657,50 @@ def detect_pipeline(f: str, op: str = 'model'): pipeline = None, None return pipeline, guess +def compile_diffusers(sd_model): + try: + if shared.opts.ipex_optimize: + import intel_extension_for_pytorch as ipex # pylint: disable=import-error, unused-import + sd_model.unet.training = False + sd_model.unet = ipex.optimize(sd_model.unet, dtype=devices.dtype_unet, inplace=True, weights_prepack=False) # pylint: disable=attribute-defined-outside-init + if hasattr(sd_model, 'vae'): + sd_model.vae.training = False + sd_model.vae = ipex.optimize(sd_model.vae, dtype=devices.dtype_vae, inplace=True, weights_prepack=False) # pylint: disable=attribute-defined-outside-init + if hasattr(sd_model, 'movq'): + sd_model.movq.training = False + sd_model.movq = ipex.optimize(sd_model.movq, dtype=devices.dtype_vae, inplace=True, weights_prepack=False) # pylint: disable=attribute-defined-outside-init + shared.log.info("Applied IPEX Optimize.") + except Exception as err: + shared.log.warning(f"IPEX Optimize not supported: {err}") + + try: + if shared.opts.cuda_compile and shared.opts.cuda_compile_backend != 'none': + shared.log.info(f"Compiling pipeline={sd_model.__class__.__name__} shape={8 * sd_model.unet.config.sample_size} mode={shared.opts.cuda_compile_backend}") + import torch._dynamo # pylint: disable=unused-import,redefined-outer-name + if shared.opts.cuda_compile_backend == "openvino_fx": + torch._dynamo.reset() # pylint: disable=protected-access + from modules.intel.openvino import openvino_fx, openvino_clear_caches # pylint: disable=unused-import + openvino_clear_caches() + torch._dynamo.eval_frame.check_if_dynamo_supported = lambda: True # pylint: disable=protected-access + if shared.compiled_model_state is None: + shared.compiled_model_state = CompiledModelState() + shared.compiled_model_state.first_pass = True if not shared.opts.cuda_compile_precompile else False + log_level = logging.WARNING if shared.opts.cuda_compile_verbose else logging.CRITICAL # pylint: disable=protected-access + if hasattr(torch, '_logging'): + torch._logging.set_logs(dynamo=log_level, aot=log_level, inductor=log_level) # pylint: disable=protected-access + torch._dynamo.config.verbose = shared.opts.cuda_compile_verbose # pylint: disable=protected-access + torch._dynamo.config.suppress_errors = shared.opts.cuda_compile_errors # pylint: disable=protected-access + sd_model.unet = torch.compile(sd_model.unet, mode=shared.opts.cuda_compile_mode, backend=shared.opts.cuda_compile_backend, fullgraph=shared.opts.cuda_compile_fullgraph) # pylint: disable=attribute-defined-outside-init + if hasattr(sd_model, 'vae'): + sd_model.vae.decode = torch.compile(sd_model.vae.decode, mode=shared.opts.cuda_compile_mode, backend=shared.opts.cuda_compile_backend, fullgraph=shared.opts.cuda_compile_fullgraph) # pylint: disable=attribute-defined-outside-init + if hasattr(sd_model, 'movq'): + sd_model.movq.decode = torch.compile(sd_model.movq.decode, mode=shared.opts.cuda_compile_mode, backend=shared.opts.cuda_compile_backend, fullgraph=shared.opts.cuda_compile_fullgraph) # pylint: disable=attribute-defined-outside-init + if shared.opts.cuda_compile_precompile: + sd_model("dummy prompt") + shared.log.info("Complilation done.") + return sd_model + except Exception as err: + shared.log.warning(f"Model compile not supported: {err}") def load_diffuser(checkpoint_info=None, already_loaded_state_dict=None, timer=None, op='model'): # pylint: disable=unused-argument import torch # pylint: disable=reimported,redefined-outer-name @@ -869,45 +913,8 @@ def load_diffuser(checkpoint_info=None, already_loaded_state_dict=None, timer=No base_sent_to_cpu=True elif not sd_model.has_accelerate: sd_model.to(devices.device) - try: - if shared.opts.ipex_optimize: - sd_model.unet.training = False - sd_model.unet = torch.xpu.optimize(sd_model.unet, dtype=devices.dtype_unet, inplace=True, weights_prepack=False) # pylint: disable=attribute-defined-outside-init - if hasattr(sd_model, 'vae'): - sd_model.vae.training = False - sd_model.vae = torch.xpu.optimize(sd_model.vae, dtype=devices.dtype_vae, inplace=True, weights_prepack=False) # pylint: disable=attribute-defined-outside-init - if hasattr(sd_model, 'movq'): - sd_model.movq.training = False - sd_model.movq = torch.xpu.optimize(sd_model.movq, dtype=devices.dtype_vae, inplace=True, weights_prepack=False) # pylint: disable=attribute-defined-outside-init - shared.log.info("Applied IPEX Optimize.") - except Exception as err: - shared.log.warning(f"IPEX Optimize not supported: {err}") - try: - if shared.opts.cuda_compile and shared.opts.cuda_compile_backend != 'none': - shared.log.info(f"Compiling pipeline={sd_model.__class__.__name__} shape={8 * sd_model.unet.config.sample_size} mode={shared.opts.cuda_compile_backend}") - import torch._dynamo # pylint: disable=unused-import,redefined-outer-name - if shared.opts.cuda_compile_backend == "openvino_fx": - torch._dynamo.reset() # pylint: disable=protected-access - from modules.intel.openvino import openvino_fx, openvino_clear_caches # pylint: disable=unused-import - openvino_clear_caches() - torch._dynamo.eval_frame.check_if_dynamo_supported = lambda: True # pylint: disable=protected-access - shared.compiled_model_state = CompiledModelState() - shared.compiled_model_state.first_pass = True if not shared.opts.cuda_compile_precompile else False - log_level = logging.WARNING if shared.opts.cuda_compile_verbose else logging.CRITICAL # pylint: disable=protected-access - if hasattr(torch, '_logging'): - torch._logging.set_logs(dynamo=log_level, aot=log_level, inductor=log_level) # pylint: disable=protected-access - torch._dynamo.config.verbose = shared.opts.cuda_compile_verbose # pylint: disable=protected-access - torch._dynamo.config.suppress_errors = shared.opts.cuda_compile_errors # pylint: disable=protected-access - sd_model.unet = torch.compile(sd_model.unet, mode=shared.opts.cuda_compile_mode, backend=shared.opts.cuda_compile_backend, fullgraph=shared.opts.cuda_compile_fullgraph) # pylint: disable=attribute-defined-outside-init - if hasattr(sd_model, 'vae'): - sd_model.vae.decode = torch.compile(sd_model.vae.decode, mode=shared.opts.cuda_compile_mode, backend=shared.opts.cuda_compile_backend, fullgraph=shared.opts.cuda_compile_fullgraph) # pylint: disable=attribute-defined-outside-init - if hasattr(sd_model, 'movq'): - sd_model.movq.decode = torch.compile(sd_model.movq.decode, mode=shared.opts.cuda_compile_mode, backend=shared.opts.cuda_compile_backend, fullgraph=shared.opts.cuda_compile_fullgraph) # pylint: disable=attribute-defined-outside-init - if shared.opts.cuda_compile_precompile: - sd_model("dummy prompt") - shared.log.info("Complilation done.") - except Exception as err: - shared.log.warning(f"Model compile not supported: {err}") + + sd_model = compile_diffusers(sd_model) if sd_model is None: shared.log.error('Diffuser model not loaded')