From 19989971891c8508eaa42db6e3f6318a3a91f1ac Mon Sep 17 00:00:00 2001 From: Disty0 Date: Tue, 31 Dec 2024 17:45:23 +0300 Subject: [PATCH] OpenVINO fix shapes resolution change and disable re-compile --- CHANGELOG.md | 3 + modules/intel/openvino/__init__.py | 103 +++++++++++++++++++---------- modules/sd_models_compile.py | 60 ++++++++++------- 3 files changed, 107 insertions(+), 59 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index d9b2b0aed..0559b184a 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -30,6 +30,9 @@ NYE refresh release with quite a few optimizatios and bug fixes... - startup optimizatios - **Torch**: - support for `torch==2.6.0` +- **OpenVINO**: + - disable re-compile on resolution change + - fix shape mismatch on resolution change - **Fixes**: - flux pipeline switches: txt/img/inpaint - flux custom unet loader for bnb diff --git a/modules/intel/openvino/__init__.py b/modules/intel/openvino/__init__.py index 157d26d96..f441e6907 100644 --- a/modules/intel/openvino/__init__.py +++ b/modules/intel/openvino/__init__.py @@ -32,16 +32,24 @@ DEFAULT_OPENVINO_PYTHON_CONFIG = MappingProxyType( class OpenVINOGraphModule(torch.nn.Module): - def __init__(self, gm, partition_id, use_python_fusion_cache, model_hash_str: str = None, file_name=""): + def __init__(self, gm, partition_id, use_python_fusion_cache, model_hash_str: str = None, file_name="", signature="", int_inputs=[]): super().__init__() self.gm = gm + self.int_inputs = int_inputs self.partition_id = partition_id + self.signature = signature self.executor_parameters = {"use_python_fusion_cache": use_python_fusion_cache, "model_hash_str": model_hash_str} self.file_name = file_name def __call__(self, *args): - result = openvino_execute(self.gm, *args, executor_parameters=self.executor_parameters, partition_id=self.partition_id, file_name=self.file_name) + ov_inputs = [] + for arg in args: + if not isinstance(arg, int): + ov_inputs.append(arg) + for idx, int_input in self.int_inputs: + ov_inputs.insert(idx, int_input) + result = openvino_execute(self.gm, *ov_inputs, executor_parameters=self.executor_parameters, partition_id=self.partition_id, file_name=self.file_name, signature=self.signature) return result @@ -111,10 +119,7 @@ def cached_model_name(model_hash_str, device, args, cache_root, reversed = False else: inputs_str += "_" + "torch.SymInt1" elif isinstance(input_data, int): - if reversed: - inputs_str = "_" + "int" + inputs_str - else: - inputs_str += "_" + "int" + pass else: if reversed: inputs_str = "_" + str(input_data.type()) + str(input_data.size())[11:-1].replace(" ", "") + inputs_str @@ -174,16 +179,13 @@ def openvino_compile(gm: GraphModule, *example_inputs, model_hash_str: str = Non input_types.append(torch.SymInt) input_shapes.append(torch.Size([1])) elif isinstance(input_data, int): - input_types.append(torch.int64) - input_shapes.append(torch.Size([1])) + pass else: input_types.append(input_data.type()) input_shapes.append(input_data.size()) decoder = TorchFXPythonDecoder(gm, input_shapes=input_shapes, input_types=input_types) - im = fe.load(decoder) - om = fe.convert(im) if file_name is not None: @@ -206,13 +208,13 @@ def openvino_compile(gm: GraphModule, *example_inputs, model_hash_str: str = Non torch.bool: Type.boolean } + idx_minus = 0 for idx, input_data in enumerate(example_inputs): if isinstance(input_data, int): - om.inputs[idx].get_node().set_element_type(dtype_mapping[torch.int64]) - om.inputs[idx].get_node().set_partial_shape(PartialShape(list(torch.Size([1])))) + idx_minus += 1 else: - 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.inputs[idx-idx_minus].get_node().set_element_type(dtype_mapping[input_data.dtype]) + om.inputs[idx-idx_minus].get_node().set_partial_shape(PartialShape(list(input_data.shape))) om.validate_nodes_and_infer_types() if shared.opts.nncf_quantize and not dont_use_quant: @@ -305,35 +307,49 @@ def openvino_compile_cached_model(cached_model_path, *example_inputs): return compiled_model -def openvino_execute(gm: GraphModule, *args, executor_parameters=None, partition_id, file_name=""): +def openvino_execute(gm: GraphModule, *args, executor_parameters=None, partition_id=None, file_name="", signature=""): + if isinstance(gm, OpenVINOGraphModule): + partition_id = gm.partition_id + file_name = gm.file_name + signature = gm.signature + executor_parameters = gm.executor_parameters + gm = gm.gm executor_parameters = executor_parameters or DEFAULT_OPENVINO_PYTHON_CONFIG - use_cache = executor_parameters.get( + if partition_id is not None: + model_cache_str = str(partition_id) + "_" + signature + elif signature: + model_cache_str = signature + + use_cache = model_cache_str and 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) if model_hash_str is not None: - model_hash_str = model_hash_str + str(partition_id) + model_hash_str = model_hash_str + str(partition_id) if partition_id is not None else "" - if use_cache and (partition_id in shared.compiled_model_state.compiled_cache): - compiled = shared.compiled_model_state.compiled_cache[partition_id] - req = shared.compiled_model_state.req_cache[partition_id] + if use_cache and (model_cache_str in shared.compiled_model_state.compiled_cache): + compiled = shared.compiled_model_state.compiled_cache[model_cache_str] + req = shared.compiled_model_state.req_cache[model_cache_str] 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) - shared.compiled_model_state.compiled_cache[partition_id] = compiled + if use_cache: + shared.compiled_model_state.compiled_cache[model_cache_str] = compiled req = compiled.create_infer_request() - shared.compiled_model_state.req_cache[partition_id] = req + if use_cache: + shared.compiled_model_state.req_cache[model_cache_str] = req flat_args, _ = tree_flatten(args) ov_inputs = [] for arg in flat_args: - ov_inputs.append((arg if isinstance(arg, int) else arg.detach().cpu().numpy())) + if not isinstance(arg, int): + ov_inputs.append((arg.detach().cpu().numpy())) res = req.infer(ov_inputs, share_inputs=True, share_outputs=True) @@ -352,33 +368,50 @@ def openvino_execute_partitioned(gm: GraphModule, *args, executor_parameters=Non ) model_hash_str = executor_parameters.get("model_hash_str", None) - signature = str(id(gm)) + signature = "signature" #str(id(gm)) + idx_minus = 0 + int_inputs = [] 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(" ", "") + if isinstance(input_data, int): + int_inputs.append([idx, input_data]) + idx_minus += 1 + elif isinstance(input_data, torch.Tensor): + signature = signature + "_" + str(idx-idx_minus) + ":" + 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) + ")" + signature = signature + "_" + str(idx-idx_minus) + ":" + type(input_data).__name__ + ":val(" + str(input_data) + ")" if signature not in shared.compiled_model_state.partitioned_modules: - shared.compiled_model_state.partitioned_modules[signature] = partition_graph(gm, use_python_fusion_cache=use_python_fusion_cache, - model_hash_str=model_hash_str, file_name=file_name) + shared.compiled_model_state.partitioned_modules[signature] = partition_graph(gm, use_python_fusion_cache=use_python_fusion_cache, + model_hash_str=model_hash_str, file_name=file_name, signature=signature, int_inputs=int_inputs) - return shared.compiled_model_state.partitioned_modules[signature](*args) + ov_inputs = [] + for arg in args: + if not isinstance(arg, int): + ov_inputs.append(arg) + for idx, int_input in shared.compiled_model_state.partitioned_modules[signature][1]: + ov_inputs.insert(idx, int_input) + return shared.compiled_model_state.partitioned_modules[signature][0](*ov_inputs) -def partition_graph(gm: GraphModule, use_python_fusion_cache: bool, model_hash_str: str = None, file_name=""): +def partition_graph(gm: GraphModule, use_python_fusion_cache: bool, model_hash_str: str = None, file_name="", signature="", int_inputs=[]): for node in gm.graph.nodes: if node.op == "call_module" and "fused_" in node.name: openvino_submodule = getattr(gm, node.name) + if isinstance(openvino_submodule, OpenVINOGraphModule): + openvino_submodule.signature = signature + int_inputs = openvino_submodule.int_inputs + if isinstance(openvino_submodule.gm, OpenVINOGraphModule): + continue 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), + OpenVINOGraphModule( + openvino_submodule, shared.compiled_model_state.partition_id, use_python_fusion_cache, + model_hash_str=model_hash_str, file_name=file_name, signature=signature, int_inputs=int_inputs), ) - shared.compiled_model_state.partition_id = shared.compiled_model_state.partition_id + 1 + shared.compiled_model_state.partition_id += 1 - return gm + return gm, int_inputs def generate_subgraph_str(tensor): diff --git a/modules/sd_models_compile.py b/modules/sd_models_compile.py index 20a7d7de2..ebc0460fe 100644 --- a/modules/sd_models_compile.py +++ b/modules/sd_models_compile.py @@ -315,9 +315,9 @@ def optimize_openvino(sd_model): shared.compiled_model_state.partitioned_modules.clear() shared.compiled_model_state = CompiledModelState() shared.compiled_model_state.is_compiled = True - shared.compiled_model_state.first_pass = True if not shared.opts.cuda_compile_precompile else False - shared.compiled_model_state.first_pass_vae = True if not shared.opts.cuda_compile_precompile else False - shared.compiled_model_state.first_pass_refiner = True if not shared.opts.cuda_compile_precompile else False + shared.compiled_model_state.first_pass = not shared.opts.cuda_compile_precompile + shared.compiled_model_state.first_pass_vae = not shared.opts.cuda_compile_precompile + shared.compiled_model_state.first_pass_refiner = not shared.opts.cuda_compile_precompile sd_models.set_accelerate(sd_model) except Exception as e: shared.log.warning(f"Model compile: task=OpenVINO: {e}") @@ -532,30 +532,42 @@ def torchao_quantization(sd_model): def openvino_recompile_model(p, hires=False, refiner=False): # recompile if a parameter changes - if 'Model' in 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 + if shared.opts.cuda_compile_backend == "openvino_fx" and 'Model' in shared.opts.cuda_compile: + 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: + openvino_first_pass = True + else: + if refiner: + openvino_first_pass = shared.compiled_model_state.first_pass_refiner + else: + openvino_first_pass = shared.compiled_model_state.first_pass + if (shared.compiled_model_state is None or + ( + not openvino_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 'Model' in shared.opts.cuda_compile and shared.opts.cuda_compile_backend == "openvino_fx": + if shared.opts.cuda_compile_backend == "openvino_fx" and 'Model' in shared.opts.cuda_compile: 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"):