Disable scheduler, tokenizer, feature extractor loading when provided (#245)

* Disable scheduler, tokenizer, feature extractor loading when given

* add test

* fix style

* fix documentation

* reduce generated number of images

Author: echarlaix

docs/source/inference.mdx

@@ -134,6 +134,7 @@ Stable Diffusion models can also be used when running inference with OpenVINO. W
 are exported to the OpenVINO format, they are decomposed into three components that are later combined during inference:
 - The text encoder
 - The U-NET
+- The VAE encoder
 - The VAE decoder
 
 Make sure you have 🤗 Diffusers installed.

optimum/intel/openvino/modeling_diffusion.py

@@ -197,7 +197,6 @@ def _from_pretrained(
         model_id = str(model_id)
         sub_models_to_load, _, _ = cls.extract_init_dict(config)
         sub_models_names = set(sub_models_to_load.keys()).intersection({"feature_extractor", "tokenizer", "scheduler"})
-        sub_models = {}
 
         if not os.path.isdir(model_id):
             patterns = set(config.keys())
@@ -231,16 +230,19 @@ def _from_pretrained(
         new_model_save_dir = Path(model_id)
 
         for name in sub_models_names:
+            # Check if the subcomponent needs to be loaded
+            if kwargs.get(name, None) is not None:
+                continue
             library_name, library_classes = sub_models_to_load[name]
             if library_classes is not None:
                 library = importlib.import_module(library_name)
                 class_obj = getattr(library, library_classes)
                 load_method = getattr(class_obj, "from_pretrained")
                 # Check if the module is in a subdirectory
                 if (new_model_save_dir / name).is_dir():
-                    sub_models[name] = load_method(new_model_save_dir / name)
+                    kwargs[name] = load_method(new_model_save_dir / name)
                 else:
-                    sub_models[name] = load_method(new_model_save_dir)
+                    kwargs[name] = load_method(new_model_save_dir)
 
         vae_decoder = cls.load_model(
             new_model_save_dir / DIFFUSION_MODEL_VAE_DECODER_SUBFOLDER / vae_decoder_file_name
@@ -260,9 +262,9 @@ def _from_pretrained(
             text_encoder=text_encoder,
             unet=unet,
             config=config,
-            tokenizer=sub_models["tokenizer"],
-            scheduler=sub_models["scheduler"],
-            feature_extractor=sub_models.pop("feature_extractor", None),
+            tokenizer=kwargs.pop("tokenizer"),
+            scheduler=kwargs.pop("scheduler"),
+            feature_extractor=kwargs.pop("feature_extractor", None),
             vae_encoder=vae_encoder,
             model_save_dir=model_save_dir,
             **kwargs,
@@ -279,6 +281,9 @@ def _from_transformers(
         cache_dir: Optional[str] = None,
         local_files_only: bool = False,
         task: Optional[str] = None,
+        tokenizer: "CLIPTokenizer" = None,
+        scheduler: Union["DDIMScheduler", "PNDMScheduler", "LMSDiscreteScheduler"] = None,
+        feature_extractor: Optional["CLIPFeatureExtractor"] = None,
         **kwargs,
     ):
         if task is None:
@@ -303,13 +308,7 @@ def _from_transformers(
             os.path.join(DIFFUSION_MODEL_VAE_DECODER_SUBFOLDER, ONNX_WEIGHTS_NAME),
         ]
         models_and_onnx_configs = get_stable_diffusion_models_for_export(model)
-
         model.save_config(save_dir_path)
-        model.tokenizer.save_pretrained(save_dir_path.joinpath("tokenizer"))
-        model.scheduler.save_pretrained(save_dir_path.joinpath("scheduler"))
-        if model.feature_extractor is not None:
-            model.feature_extractor.save_pretrained(save_dir_path.joinpath("feature_extractor"))
-
         export_models(
             models_and_onnx_configs=models_and_onnx_configs,
             output_dir=save_dir_path,
@@ -325,7 +324,10 @@ def _from_transformers(
             force_download=force_download,
             cache_dir=cache_dir,
             local_files_only=local_files_only,
-            model_save_dir=save_dir,  # important
+            model_save_dir=save_dir,
+            tokenizer=tokenizer or model.tokenizer,
+            scheduler=scheduler or model.scheduler,
+            feature_extractor=feature_extractor or model.feature_extractor,
             **kwargs,
         )
 

setup.py

@@ -29,14 +29,14 @@
         "neural-compressor>=2.0.0",
         "onnx",
         "onnxruntime",
-        "torch<2.0.0", # remove after neural-compressor next release
+        "torch<2.0.0",  # remove after neural-compressor next release
         "intel-extension-for-pytorch<2.0.0",
     ],
     "openvino": [
         "openvino>=2023.0.0.dev20230217",
         "onnx",
         "onnxruntime",
-        "torch<2.0.0", # remove after optimum next release
+        "torch<2.0.0",  # remove after optimum next release
     ],
     "nncf": ["nncf>=2.4.0", "openvino-dev>=2023.0.0.dev20230217"],
     "ipex": ["intel-extension-for-pytorch"],

tests/openvino/test_modeling.py

@@ -114,7 +114,7 @@ def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
         self.OV_MODEL_ID = "echarlaix/distilbert-base-uncased-finetuned-sst-2-english-openvino"
         self.OV_SEQ2SEQ_MODEL_ID = "echarlaix/t5-small-openvino"
-        self.OV_STABLE_DIFFUSION_MODEL_ID = "hf-internal-testing/tiny-stable-diffusion-openvino"
+        self.OV_DIFFUSION_MODEL_ID = "hf-internal-testing/tiny-stable-diffusion-openvino"
 
     def test_load_from_hub_and_save_model(self):
         tokenizer = AutoTokenizer.from_pretrained(self.OV_MODEL_ID)
@@ -152,8 +152,9 @@ def test_load_from_hub_and_save_seq2seq_model(self):
         outputs = model.generate(**tokens)
         self.assertTrue(torch.equal(loaded_model_outputs, outputs))
 
+    @require_diffusers
     def test_load_from_hub_and_save_stable_diffusion_model(self):
-        loaded_pipeline = OVStableDiffusionPipeline.from_pretrained(self.OV_STABLE_DIFFUSION_MODEL_ID, compile=False)
+        loaded_pipeline = OVStableDiffusionPipeline.from_pretrained(self.OV_DIFFUSION_MODEL_ID, compile=False)
         self.assertIsInstance(loaded_pipeline.config, Dict)
         prompt = "sailing ship in storm by Leonardo da Vinci"
         height = 16
@@ -704,23 +705,23 @@ def test_compare_to_diffusers(self, model_arch: str):
     @parameterized.expand(SUPPORTED_ARCHITECTURES)
     @require_diffusers
     def test_num_images_per_prompt(self, model_arch: str):
+        from diffusers import DPMSolverMultistepScheduler
+
         model_id = MODEL_NAMES[model_arch]
-        num_images_per_prompt = 4
-        batch_size = 6
-        pipeline = OVStableDiffusionPipeline.from_pretrained(model_id, export=True)
+        scheduler = DPMSolverMultistepScheduler.from_pretrained(model_id, subfolder="scheduler")
+        pipeline = OVStableDiffusionPipeline.from_pretrained(model_id, export=True, scheduler=scheduler)
         prompt = "sailing ship in storm by Leonardo da Vinci"
-        outputs = pipeline(prompt, num_inference_steps=2, output_type="np").images
-        self.assertEqual(outputs.shape, (1, 128, 128, 3))
-        outputs = pipeline(
-            prompt, num_inference_steps=2, num_images_per_prompt=num_images_per_prompt, output_type="np"
-        ).images
-        self.assertEqual(outputs.shape, (num_images_per_prompt, 128, 128, 3))
-        outputs = pipeline([prompt] * batch_size, num_inference_steps=2, output_type="np").images
-        self.assertEqual(outputs.shape, (batch_size, 128, 128, 3))
+
+        for batch_size in [1, 3]:
+            for num_images in [1, 2]:
+                outputs = pipeline(
+                    [prompt] * batch_size, num_inference_steps=2, num_images_per_prompt=num_images, output_type="np"
+                )
+                self.assertEqual(outputs.images.shape, (batch_size * num_images, 128, 128, 3))
 
     @parameterized.expand(SUPPORTED_ARCHITECTURES)
     @require_diffusers
-    def test_num_images_per_prompt(self, model_arch: str):
+    def test_num_images_per_prompt_static_model(self, model_arch: str):
         model_id = MODEL_NAMES[model_arch]
         batch_size = 3
         num_images_per_prompt = 4