fix loading safety checker in pipeline

huggingface · Nov 4, 2024 · 591dbfa · 591dbfa
1 parent 74d2161
commit 591dbfa
Showing 1 changed file with 8 additions and 3 deletions.
diff --git 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)