Load diffusers in native FP16/BF16 precision to reduce the memory usage (#1033)

* Load diffusers in native FP16/BF16 precision to reduce the memory usage

* Apply suggestions from code review

Co-authored-by: Ekaterina Aidova <ekaterina.aidova@intel.com>

* Fix code

* Update optimum/exporters/openvino/__main__.py

Co-authored-by: Roman Kazantsev <roman.kazantsev@intel.com>

* Find first floating point tensor instead of first tensor

* Fix style

* Update optimum/exporters/openvino/__main__.py

* Update optimum/exporters/openvino/__main__.py

* Check if safetensors available

* Fix style

* Extract code in the utils fuction

---------

Co-authored-by: Ekaterina Aidova <ekaterina.aidova@intel.com>
Co-authored-by: Roman Kazantsev <roman.kazantsev@intel.com>

Author: 3 people

optimum/exporters/openvino/__main__.py

@@ -41,7 +41,12 @@
 )
 from optimum.utils.save_utils import maybe_load_preprocessors
 
-from .utils import _MAX_UNCOMPRESSED_SIZE, MULTI_MODAL_TEXT_GENERATION_MODELS, clear_class_registry
+from .utils import (
+    _MAX_UNCOMPRESSED_SIZE,
+    MULTI_MODAL_TEXT_GENERATION_MODELS,
+    clear_class_registry,
+    deduce_diffusers_dtype,
+)
 
 
 FORCE_ATTN_MODEL_CLASSES = {"phi3-v": "eager"}
@@ -332,6 +337,19 @@ class StoreAttr(object):
                 return model
 
             GPTQQuantizer.post_init_model = post_init_model
+    elif library_name == "diffusers" and is_openvino_version(">=", "2024.6"):
+        dtype = deduce_diffusers_dtype(
+            model_name_or_path,
+            revision=revision,
+            cache_dir=cache_dir,
+            token=token,
+            local_files_only=local_files_only,
+            force_download=force_download,
+            trust_remote_code=trust_remote_code,
+        )
+        if dtype in [torch.float16, torch.bfloat16]:
+            loading_kwargs["torch_dtype"] = dtype
+            patch_16bit = True
 
     if library_name == "open_clip":
         model = _OpenClipForZeroShotImageClassification.from_pretrained(model_name_or_path, cache_dir=cache_dir)

optimum/exporters/openvino/utils.py

@@ -26,6 +26,7 @@
 from optimum.exporters import TasksManager
 from optimum.exporters.onnx.base import OnnxConfig
 from optimum.intel.utils import is_transformers_version
+from optimum.intel.utils.import_utils import is_safetensors_available
 from optimum.utils import is_diffusers_available
 from optimum.utils.save_utils import maybe_save_preprocessors
 
@@ -232,6 +233,41 @@ def save_config(config, save_dir):
         config.to_json_file(output_config_file, use_diff=True)
 
 
+def deduce_diffusers_dtype(model_name_or_path, **loading_kwargs):
+    dtype = None
+    if is_safetensors_available():
+        if Path(model_name_or_path).is_dir():
+            path = Path(model_name_or_path)
+        else:
+            from diffusers import DiffusionPipeline
+
+            path = DiffusionPipeline.download(model_name_or_path, **loading_kwargs)
+        model_part_name = None
+        if (path / "transformer").is_dir():
+            model_part_name = "transformer"
+        elif (path / "unet").is_dir():
+            model_part_name = "unet"
+        if model_part_name:
+            directory = path / model_part_name
+            safetensors_files = [
+                filename for filename in directory.glob("*.safetensors") if len(filename.suffixes) == 1
+            ]
+            safetensors_file = None
+            if len(safetensors_files) > 0:
+                safetensors_file = safetensors_files.pop(0)
+            if safetensors_file:
+                from safetensors import safe_open
+
+                with safe_open(safetensors_file, framework="pt", device="cpu") as f:
+                    if len(f.keys()) > 0:
+                        for key in f.keys():
+                            tensor = f.get_tensor(key)
+                            if tensor.dtype.is_floating_point:
+                                dtype = tensor.dtype
+                                break
+    return dtype
+
+
 def save_preprocessors(
     preprocessors: List, config: PretrainedConfig, output: Union[str, Path], trust_remote_code: bool
 ):