Export hybrid StableDiffusion models via optimum-cli

Author: l-bat

optimum/commands/export/openvino.py

@@ -104,6 +104,16 @@ def parse_args_openvino(parser: "ArgumentParser"):
         default=None,
         help=("The group size to use for quantization. Recommended value is 128 and -1 uses per-column quantization."),
     )
+    optional_group.add_argument(
+        "--dataset",
+        type=str,
+        default=None,
+        help=(
+            "The dataset used for data-aware compression or quantization with NNCF. "
+            "You can use the one from the list ['wikitext2','c4','c4-new','ptb','ptb-new'] for LLLMs "
+            "or ['conceptual_captions','laion/220k-GPT4Vision-captions-from-LIVIS','laion/filtered-wit'] for diffusion models."
+        ),
+    )
     optional_group.add_argument(
         "--disable-stateful",
         action="store_true",
@@ -187,6 +197,7 @@ def run(self):
                 )
                 quantization_config["sym"] = "asym" not in self.args.weight_format
                 quantization_config["group_size"] = 128 if "128" in self.args.weight_format else 64
+            quantization_config["dataset"] = self.args.dataset
             ov_config = OVConfig(quantization_config=quantization_config)
 
         # TODO : add input shapes

optimum/exporters/openvino/__main__.py

@@ -263,6 +263,44 @@ class StoreAttr(object):
         **loading_kwargs,
     )
 
+    if (
+        library_name == "diffusers"
+        and ov_config
+        and ov_config.quantization_config
+        and "dataset" in ov_config.quantization_config
+    ):
+        import huggingface_hub
+
+        model_info = huggingface_hub.model_info(model_name_or_path, revision=revision)
+        class_name = model_info.config["diffusers"]["_class_name"]
+        model_cls = None
+        if class_name == "LatentConsistencyModelPipeline":
+            from optimum.intel import OVLatentConsistencyModelPipeline
+
+            model_cls = OVLatentConsistencyModelPipeline
+        elif class_name == "StableDiffusionXLPipeline":
+            from optimum.intel import OVStableDiffusionXLPipeline
+
+            model_cls = OVStableDiffusionXLPipeline
+        elif class_name == "StableDiffusionPipeline":
+            from optimum.intel import OVStableDiffusionPipeline
+
+            model_cls = OVStableDiffusionPipeline
+
+        if model_cls:
+            model = model_cls.from_pretrained(
+                model_id=model_name_or_path,
+                export=True,
+                quantization_config=ov_config.quantization_config,
+                cache_dir=cache_dir,
+                trust_remote_code=trust_remote_code,
+                revision=revision,
+                force_download=force_download,
+                use_auth_token=use_auth_token,
+            )
+            model.save_pretrained(output)
+            return
+
     needs_pad_token_id = task == "text-classification" and getattr(model.config, "pad_token_id", None) is None
 
     if needs_pad_token_id:

optimum/intel/openvino/modeling_diffusion.py

@@ -385,7 +385,7 @@ def transform_fn(data_item):
                 self.__call__(**inputs, height=height, width=width)
             else:
                 self.__call__(*inputs, height=height, width=width)
-            if len(calibration_data) > num_samples:
+            if len(calibration_data) >= num_samples:
                 break
 
         self.unet.request = self.unet.request.request