Fix export for diffusion models for image-to-image and inpainting task (

#848)

* fix export feature for diffusion models

* set test to slow

* fix test

optimum/intel/openvino/modeling_diffusion.py

@@ -77,7 +77,7 @@
 class OVStableDiffusionPipelineBase(OVBaseModel, OVTextualInversionLoaderMixin):
     auto_model_class = StableDiffusionPipeline
     config_name = "model_index.json"
-    export_feature = "stable-diffusion"
+    export_feature = "text-to-image"
 
     def __init__(
         self,
@@ -799,6 +799,8 @@ def __call__(
 
 
 class OVStableDiffusionImg2ImgPipeline(OVStableDiffusionPipelineBase, StableDiffusionImg2ImgPipelineMixin):
+    export_feature = "image-to-image"
+
     def __call__(
         self,
         prompt: Optional[Union[str, List[str]]] = None,
@@ -841,6 +843,8 @@ def __call__(
 
 
 class OVStableDiffusionInpaintPipeline(OVStableDiffusionPipelineBase, StableDiffusionInpaintPipelineMixin):
+    export_feature = "inpainting"
+
     def __call__(
         self,
         prompt: Optional[Union[str, List[str]]],
@@ -912,7 +916,6 @@ def __call__(
 
 class OVStableDiffusionXLPipelineBase(OVStableDiffusionPipelineBase):
     auto_model_class = StableDiffusionXLPipeline
-    export_feature = "stable-diffusion-xl"
 
     def __init__(self, *args, add_watermarker: Optional[bool] = None, **kwargs):
         super().__init__(*args, **kwargs)
@@ -988,6 +991,7 @@ def __call__(
 
 class OVStableDiffusionXLImg2ImgPipeline(OVStableDiffusionXLPipelineBase, StableDiffusionXLImg2ImgPipelineMixin):
     auto_model_class = StableDiffusionXLImg2ImgPipeline
+    export_feature = "image-to-image"
 
     def __call__(
         self,

tests/openvino/test_modeling.py

@@ -249,13 +249,15 @@ def test_load_from_hub_and_save_stable_diffusion_model(self):
         del pipeline
         gc.collect()
 
+    @pytest.mark.run_slow
+    @slow
     def test_load_model_from_hub_private_with_token(self):
         token = os.environ.get("HF_HUB_READ_TOKEN", None)
         if token is None:
             self.skipTest("Test requires a token `HF_HUB_READ_TOKEN` in the environment variable")
 
         model = OVModelForCausalLM.from_pretrained(
-            "optimum-internal-testing/tiny-random-phi-private", use_auth_token=token, revision="openvino"
+            "optimum-internal-testing/tiny-random-phi-private", token=token, revision="openvino"
         )
         self.assertIsInstance(model.config, PretrainedConfig)
 

tests/openvino/test_quantization.py

@@ -395,9 +395,9 @@ def test_ovmodel_load_with_compressed_weights(self, model_cls, model_type):
 
         if model.export_feature.startswith("text2text-generation"):
             models = [model.encoder, model.decoder, model.decoder_with_past]
-        elif model.export_feature.startswith("stable-diffusion"):
+        elif model.export_feature == "text-to-image":
             models = [model.unet, model.vae_encoder, model.vae_decoder]
-            models.append(model.text_encoder if model.export_feature == "stable-diffusion" else model.text_encoder_2)
+            models.append(model.text_encoder if model_type == "stable-diffusion" else model.text_encoder_2)
         else:
             models = [model]
 
@@ -500,12 +500,11 @@ def test_ovmodel_stateful_load_with_compressed_weights(self, model_cls, model_ty
     @parameterized.expand(SUPPORTED_ARCHITECTURES_WITH_AUTO_COMPRESSION)
     def test_ovmodel_load_with_uncompressed_weights(self, model_cls, model_type):
         model = model_cls.from_pretrained(MODEL_NAMES[model_type], export=True, load_in_8bit=False)
-
         if model.export_feature.startswith("text2text-generation"):
             models = [model.encoder, model.decoder, model.decoder_with_past]
-        elif model.export_feature.startswith("stable-diffusion"):
+        elif model.export_feature == "text-to-image":
             models = [model.unet, model.vae_encoder, model.vae_decoder]
-            models.append(model.text_encoder if model.export_feature == "stable-diffusion" else model.text_encoder_2)
+            models.append(model.text_encoder if model_type == "stable-diffusion" else model.text_encoder_2)
         else:
             models = [model]