OpenVINO fix shapes resolution change and disable re-compile

This commit is contained in:
Disty0
2024-12-31 17:45:23 +03:00
parent cd32b966b2
commit 1998997189
3 changed files with 107 additions and 59 deletions
+3
View File
@@ -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
+68 -35
View File
@@ -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):
+36 -24
View File
@@ -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"):