OpenVINO fix caching and recompile when using Lora

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

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()

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')

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')