update tests

Author: eaidova

optimum/commands/export/openvino.py

@@ -416,6 +416,10 @@ def run(self):
                 from optimum.intel import OVFluxPipeline
 
                 model_cls = OVFluxPipeline
+            elif class_name == "SanaPipeline":
+                from optimum.intel import OVSanaPipeline
+
+                model_cls = OVSanaPipeline
             else:
                 raise NotImplementedError(f"Quantization in hybrid mode isn't supported for class {class_name}.")
 

optimum/exporters/openvino/convert.py

@@ -1047,7 +1047,6 @@ def get_diffusion_models_for_export_ext(
             sd3_pipes.append(StableDiffusion3InpaintPipeline)
 
         is_sd3 = isinstance(pipeline, tuple(sd3_pipes))
-        logger.warn(f"IS SD3 {pipeline} {is_sd3}")
     else:
         is_sd3 = False
 
@@ -1185,7 +1184,6 @@ def get_sd3_models_for_export(pipeline, exporter, int_dtype, float_dtype):
         task="semantic-segmentation",
         model_type="sd3-transformer",
     )
-    logger.warn(f"TRANSFORMER COFG {export_config_constructor}")
     transformer_export_config = export_config_constructor(
         pipeline.transformer.config, int_dtype=int_dtype, float_dtype=float_dtype
     )

optimum/intel/openvino/modeling_diffusion.py

@@ -844,17 +844,23 @@ def reshape(
 
         if self.text_encoder is not None:
             self.text_encoder.model = self._reshape_text_encoder(
-                self.text_encoder.model, batch_size, self.tokenizer.model_max_length
+                self.text_encoder.model,
+                batch_size,
+                self.tokenizer.model_max_length if "Gemma" not in self.tokenizer.__class__.__name__ else -1,
             )
 
         if self.text_encoder_2 is not None:
             self.text_encoder_2.model = self._reshape_text_encoder(
-                self.text_encoder_2.model, batch_size, self.tokenizer_2.model_max_length
+                self.text_encoder_2.model,
+                batch_size,
+                self.tokenizer_2.model_max_length if "Gemma" not in self.tokenizer.__class__.__name__ else -1,
             )
 
         if self.text_encoder_3 is not None:
             self.text_encoder_3.model = self._reshape_text_encoder(
-                self.text_encoder_3.model, batch_size, self.tokenizer_3.model_max_length
+                self.text_encoder_3.model,
+                batch_size,
+                self.tokenizer_3.model_max_length if "Gemma" not in self.tokenizer.__class__.__name__ else -1,
             )
 
         self.clear_requests()

optimum/intel/openvino/utils.py

@@ -125,6 +125,7 @@
     "stable-diffusion": "OVStableDiffusionPipeline",
     "stable-diffusion-xl": "OVStableDiffusionXLPipeline",
     "stable-diffusion-3": "OVStableDiffusion3Pipeline",
+    "sana": "OVSanaPipeline",
     "flux": "OVFluxPipeline",
     "flux-fill": "OVFluxFillPipeline",
     "pix2struct": "OVModelForPix2Struct",

tests/openvino/test_diffusion.py

@@ -78,7 +78,7 @@ class OVPipelineForText2ImageTest(unittest.TestCase):
     NEGATIVE_PROMPT_SUPPORT_ARCHITECTURES = ["stable-diffusion", "stable-diffusion-xl", "latent-consistency"]
     if is_transformers_version(">=", "4.40.0"):
         SUPPORTED_ARCHITECTURES.extend(["stable-diffusion-3", "flux", "sana"])
-        NEGATIVE_PROMPT_SUPPORT_ARCHITECTURES.extend(["stable-diffusion-3", "sana"])
+        NEGATIVE_PROMPT_SUPPORT_ARCHITECTURES.append(["stable-diffusion-3"])
     CALLBACK_SUPPORT_ARCHITECTURES = ["stable-diffusion", "stable-diffusion-xl", "latent-consistency"]
 
     OVMODEL_CLASS = OVPipelineForText2Image
@@ -215,6 +215,8 @@ def test_shape(self, model_arch: str):
 
         height, width, batch_size = 128, 64, 1
         inputs = self.generate_inputs(height=height, width=width, batch_size=batch_size)
+        if model_arch == "sana":
+            inputs["use_resolution_binning"] = False
 
         for output_type in ["pil", "np", "pt", "latent"]:
             inputs["output_type"] = output_type

tests/openvino/test_exporters_cli.py

@@ -44,6 +44,7 @@
     OVModelOpenCLIPForZeroShotImageClassification,
     OVModelOpenCLIPText,
     OVModelOpenCLIPVisual,
+    OVSanaPipeline,
     OVSentenceTransformer,
     OVStableDiffusion3Pipeline,
     OVStableDiffusionPipeline,
@@ -107,6 +108,7 @@ class OVCLIExportTestCase(unittest.TestCase):
         "flux": 4 if is_tokenizers_version("<", "0.20") or is_openvino_version(">=", "2024.5") else 0,
         "flux-fill": 4 if is_tokenizers_version("<", "0.20") or is_openvino_version(">=", "2024.5") else 0,
         "llava": 2 if is_tokenizers_version("<", "0.20") or is_openvino_version(">=", "2024.5") else 0,
+        "sana": 2 if is_tokenizers_version("<", "0.20.0") or is_openvino_version(">=", "2024.5") else 0,
     }
 
     SUPPORTED_SD_HYBRID_ARCHITECTURES = [
@@ -118,7 +120,7 @@ class OVCLIExportTestCase(unittest.TestCase):
     if is_transformers_version(">=", "4.45"):
         SUPPORTED_SD_HYBRID_ARCHITECTURES.append(("stable-diffusion-3", 9, 65))
         SUPPORTED_SD_HYBRID_ARCHITECTURES.append(("flux", 7, 56))
-        SUPPORTED_SD_HYBRID_ARCHITECTURES.append(("sana", 7, 56))
+        SUPPORTED_SD_HYBRID_ARCHITECTURES.append(("sana", 19, 53))
 
     SUPPORTED_QUANTIZATION_ARCHITECTURES = [
         (
@@ -357,9 +359,15 @@ def test_exporters_cli_int8(self, task: str, model_type: str):
                 models = [model.encoder, model.decoder]
                 if task.endswith("with-past") and not model.decoder.stateful:
                     models.append(model.decoder_with_past)
-            elif model_type.startswith("stable-diffusion") or model_type.startswith("flux"):
+            elif (
+                model_type.startswith("stable-diffusion")
+                or model_type.startswith("flux")
+                or model_type.startswith("sana")
+            ):
                 models = [model.unet or model.transformer, model.vae_encoder, model.vae_decoder]
-                models.append(model.text_encoder if model_type == "stable-diffusion" else model.text_encoder_2)
+                models.append(
+                    model.text_encoder if model_type in ["stable-diffusion", "sana"] else model.text_encoder_2
+                )
             elif task.startswith("image-text-to-text"):
                 models = [model.language_model, model.vision_embeddings]
             else:

tests/openvino/test_quantization.py

@@ -59,6 +59,7 @@
     OVStableDiffusionXLPipeline,
     OVStableDiffusion3Pipeline,
     OVQuantizer,
+    OVSanaPipeline,
     OVTrainer,
     OVQuantizationConfig,
     OVWeightQuantizationConfig,
@@ -543,6 +544,7 @@ class OVWeightCompressionTest(unittest.TestCase):
             [
                 (OVStableDiffusion3Pipeline, "stable-diffusion-3", 9, 65),
                 (OVFluxPipeline, "flux", 7, 56),
+                (OVSanaPipeline, "sana", 19, 53),
             ]
         )
 

tests/openvino/utils_tests.py

@@ -168,7 +168,7 @@
     "open-clip-ov": "zofinka/tiny-open-clip-model",
     "st-bert": "sentence-transformers/all-MiniLM-L6-v2",
     "st-mpnet": "sentence-transformers/all-mpnet-base-v2",
-    "sana": "/home/ea/work/my_optimum_intel/optimum-intel/tiny-random-sana",
+    "sana": "katuni4ka/tiny-random-sana",
 }
 
 
@@ -201,7 +201,7 @@
     "minicpmv": (30, 26, 1, 6),
     "nanollava": (30, 15, 1),
     "qwen2_vl": (30, 1, 1, 10),
-    "sana": (242, 34, 42, 64),
+    "sana": (58, 28, 28, 18),
 }
 
 TEST_IMAGE_URL = "http://images.cocodataset.org/val2017/000000039769.jpg"