From 591dbfaeb07dbc85f6127c853b1533ad3dad3e5a Mon Sep 17 00:00:00 2001 From: eaidova Date: Mon, 4 Nov 2024 14:36:32 +0400 Subject: [PATCH] fix loading safety checker in pipeline --- optimum/intel/openvino/modeling_diffusion.py | 11 ++++++++--- 1 file changed, 8 insertions(+), 3 deletions(-) diff --git a/optimum/intel/openvino/modeling_diffusion.py b/optimum/intel/openvino/modeling_diffusion.py index 18d8a75061..89c240e8e8 100644 --- a/optimum/intel/openvino/modeling_diffusion.py +++ b/optimum/intel/openvino/modeling_diffusion.py @@ -409,9 +409,9 @@ def _from_pretrained( "tokenizer_2": None, "tokenizer_3": None, "feature_extractor": None, - "image_encoder": None, - "safety_checker": None, } + + additional_submodels = ["image_encoder", "safety_checker"] for name in submodels.keys(): if kwargs.get(name) is not None: submodels[name] = kwargs.pop(name) @@ -437,7 +437,12 @@ def _from_pretrained( } for config_key, value in config.items(): - if config_key not in models and config_key not in kwargs and config_key not in submodels: + if ( + config_key not in models + and config_key not in kwargs + and config_key not in submodels + and config_key not in additional_submodels + ): kwargs[config_key] = value compile_only = kwargs.get("compile_only", False)