diff --git a/optimum/intel/openvino/modeling_diffusion.py b/optimum/intel/openvino/modeling_diffusion.py index f6c4fc37a8..3fa6f886b9 100644 --- a/optimum/intel/openvino/modeling_diffusion.py +++ b/optimum/intel/openvino/modeling_diffusion.py @@ -294,8 +294,12 @@ def _save_pretrained(self, save_directory: Union[str, Path]): dst_path = save_path / OV_XML_FILE_NAME dst_path.parent.mkdir(parents=True, exist_ok=True) openvino.save_model(model.model, dst_path, compress_to_fp16=False) - model_dir = model.config.get("_name_or_path", None) or model.model_save_dir - config_path = Path(model_dir) / CONFIG_NAME + model_dir = ( + self.model_save_dir + if not isinstance(self.model_save_dir, TemporaryDirectory) + else self.model_save_dir.name + ) + config_path = Path(model_dir) / save_path.name / CONFIG_NAME if config_path.is_file(): config_save_path = save_path / CONFIG_NAME shutil.copyfile(config_path, config_save_path) diff --git a/tests/openvino/test_diffusion.py b/tests/openvino/test_diffusion.py index 3e139b72a7..b302229cd7 100644 --- a/tests/openvino/test_diffusion.py +++ b/tests/openvino/test_diffusion.py @@ -359,6 +359,13 @@ def test_load_and_save_pipeline_with_safety_checker(self): self.assertTrue(model_lib in ["diffusers", "transformers"]) self.assertFalse(model_class.startswith("OV")) loaded_pipeline = self.OVMODEL_CLASS.from_pretrained(tmpdirname) + for component in ["text_encoder", "unet", "vae_encoder", "vae_decoder"]: + config = getattr(getattr(ov_pipeline, component), "config", None) + if config is not None: + loaded_config = getattr(getattr(loaded_pipeline, component), "config") + self.assertDictEqual( + config, loaded_config, f"Expected config:\n{config}\nLoaded config:|n{loaded_config}" + ) self.assertTrue(loaded_pipeline.safety_checker is not None) self.assertIsInstance(loaded_pipeline.safety_checker, StableDiffusionSafetyChecker) del loaded_pipeline