Merge branch 'itikhono/update_sdpa2pa_for_qwen' of https://github.com…

…/itikhono/openvino into itikhono/update_sdpa2pa_for_qwen

src/bindings/python/src/openvino/_ov_api.py

@@ -2,7 +2,8 @@
 # Copyright (C) 2018-2024 Intel Corporation
 # SPDX-License-Identifier: Apache-2.0
 
-from typing import Any, Iterable, Union, Optional, Dict
+from types import TracebackType
+from typing import Any, Iterable, Union, Optional, Dict, Type
 from pathlib import Path
 
 
@@ -21,30 +22,48 @@
 )
 
 
-class Model(ModelBase):
+class Model:
     def __init__(self, *args: Any, **kwargs: Any) -> None:
         if args and not kwargs:
             if isinstance(args[0], ModelBase):
-                super().__init__(args[0])
+                self.__model = ModelBase(args[0])
             elif isinstance(args[0], Node):
-                super().__init__(*args)
+                self.__model = ModelBase(*args)
             else:
-                super().__init__(*args)
+                self.__model = ModelBase(*args)
         if args and kwargs:
-            super().__init__(*args, **kwargs)
+            self.__model = ModelBase(*args, **kwargs)
         if kwargs and not args:
-            super().__init__(**kwargs)
+            self.__model = ModelBase(**kwargs)
+
+    def __getattr__(self, name: str) -> Any:
+        if self.__model is None:
+            raise AttributeError(f"'Model' object has no attribute '{name}' or attribute is no longer accessible.")
+        return getattr(self.__model, name)
 
     def clone(self) -> "Model":
-        return Model(super().clone())
+        return Model(self.__model.clone())
+
+    def __copy__(self) -> "Model":
+        raise TypeError("Cannot copy 'openvino.runtime.Model'. Please, use deepcopy instead.")
 
     def __deepcopy__(self, memo: Dict) -> "Model":
         """Returns a deepcopy of Model.
 
         :return: A copy of Model.
         :rtype: openvino.runtime.Model
         """
-        return Model(super().clone())
+        return Model(self.__model.clone())
+
+    def __enter__(self) -> "Model":
+        return self
+
+    def __exit__(self, exc_type: Type[BaseException], exc_value: BaseException, traceback: TracebackType) -> None:
+        del self.__model
+        self.__model = None
+
+    def __repr__(self) -> str:
+        return self.__model.__repr__()
 
 
 class InferRequest(_InferRequestWrapper):
@@ -500,6 +519,8 @@ def read_model(
         config: Optional[dict] = None
     ) -> Model:
         config = {} if config is None else config
+        if isinstance(model, Model):
+            model = model._Model__model
 
         if isinstance(weights, Tensor):
             return Model(super().read_model(model, weights))
@@ -543,6 +564,8 @@ def compile_model(
         :return: A compiled model.
         :rtype: openvino.runtime.CompiledModel
         """
+        if isinstance(model, Model):
+            model = model._Model__model
         if weights is None:
             if device_name is None:
                 return CompiledModel(
@@ -562,6 +585,16 @@ def compile_model(
                 weights=weights,
             )
 
+    def query_model(
+            self,
+            model: Model,
+            device_name: str,
+            config: Optional[dict] = None,
+    ) -> dict:
+        return super().query_model(model._Model__model,
+                                   device_name,
+                                   {} if config is None else config, )
+
     def import_model(
         self,
         model_stream: bytes,
@@ -637,4 +670,6 @@ def compile_model(
 
     """
     core = Core()
+    if isinstance(model, Model):
+        model = model._Model__model
     return core.compile_model(model, device_name, {} if config is None else config)

src/bindings/python/src/openvino/properties/_properties.py

@@ -16,6 +16,9 @@ def __new__(cls, prop: Callable[..., Any]):  # type: ignore
 
     def __call__(self, *args: Any) -> Callable[..., Any]:
         if args is not None:
+            from openvino import Model
+            if args and isinstance(args[0], Model):
+                return self.prop(args[0]._Model__model)
             return self.prop(*args)
         return self.prop()
 

src/bindings/python/src/openvino/test_utils/__init__.py

@@ -2,4 +2,4 @@
 # Copyright (C) 2018-2024 Intel Corporation
 # SPDX-License-Identifier: Apache-2.0
 
-from .test_utils_api import compare_functions
+from .test_api import compare_functions

src/bindings/python/src/openvino/test_utils/test_api.py

@@ -0,0 +1,10 @@
+# -*- coding: utf-8 -*-
+# Copyright (C) 2018-2024 Intel Corporation
+# SPDX-License-Identifier: Apache-2.0
+
+from .test_utils_api import compare_functions as compare_functions_base
+from openvino.runtime import Model
+
+
+def compare_functions(lhs: Model, rhs: Model, compare_tensor_names: bool = True) -> tuple:
+    return compare_functions_base(lhs._Model__model, rhs._Model__model, compare_tensor_names)

src/bindings/python/src/pyopenvino/core/offline_transformations.cpp

@@ -23,6 +23,7 @@
 
 #include "openvino/pass/low_latency.hpp"
 #include "openvino/pass/manager.hpp"
+#include "pyopenvino/utils/utils.hpp"
 
 namespace py = pybind11;
 
@@ -34,7 +35,8 @@ void regmodule_offline_transformations(py::module m) {
 
     m_offline_transformations.def(
         "apply_moc_transformations",
-        [](std::shared_ptr<ov::Model> model, bool cf, bool smart_reshape) {
+        [](py::object& ie_api_model, bool cf, bool smart_reshape) {
+            const auto model = Common::utils::convert_to_model(ie_api_model);
             ov::pass::Manager manager;
             if (smart_reshape)
                 manager.register_pass<ov::pass::SmartReshape>();
@@ -48,7 +50,8 @@ void regmodule_offline_transformations(py::module m) {
 
     m_offline_transformations.def(
         "apply_moc_legacy_transformations",
-        [](std::shared_ptr<ov::Model> model, const std::vector<std::string>& params_with_custom_types) {
+        [](py::object& ie_api_model, const std::vector<std::string>& params_with_custom_types) {
+            const auto model = Common::utils::convert_to_model(ie_api_model);
             ov::pass::Manager manager;
             manager.register_pass<ov::pass::MOCLegacyTransformations>(params_with_custom_types);
             manager.run_passes(model);
@@ -58,7 +61,8 @@ void regmodule_offline_transformations(py::module m) {
 
     m_offline_transformations.def(
         "apply_low_latency_transformation",
-        [](std::shared_ptr<ov::Model> model, bool use_const_initializer = true) {
+        [](py::object& ie_api_model, bool use_const_initializer = true) {
+            const auto model = Common::utils::convert_to_model(ie_api_model);
             ov::pass::Manager manager;
             manager.register_pass<ov::pass::LowLatency2>(use_const_initializer);
             manager.run_passes(model);
@@ -68,7 +72,8 @@ void regmodule_offline_transformations(py::module m) {
 
     m_offline_transformations.def(
         "apply_pruning_transformation",
-        [](std::shared_ptr<ov::Model> model) {
+        [](py::object& ie_api_model) {
+            const auto model = Common::utils::convert_to_model(ie_api_model);
             ov::pass::Manager manager;
             manager.register_pass<ov::pass::Pruning>();
             manager.run_passes(model);
@@ -77,7 +82,8 @@ void regmodule_offline_transformations(py::module m) {
 
     m_offline_transformations.def(
         "apply_make_stateful_transformation",
-        [](std::shared_ptr<ov::Model> model, const std::map<std::string, std::string>& param_res_names) {
+        [](py::object& ie_api_model, const std::map<std::string, std::string>& param_res_names) {
+            const auto model = Common::utils::convert_to_model(ie_api_model);
             ov::pass::Manager manager;
             manager.register_pass<ov::pass::MakeStateful>(param_res_names);
             manager.run_passes(model);
@@ -87,7 +93,8 @@ void regmodule_offline_transformations(py::module m) {
 
     m_offline_transformations.def(
         "apply_make_stateful_transformation",
-        [](std::shared_ptr<ov::Model> model, const ov::pass::MakeStateful::ParamResPairs& pairs_to_replace) {
+        [](py::object& ie_api_model, const ov::pass::MakeStateful::ParamResPairs& pairs_to_replace) {
+            const auto model = Common::utils::convert_to_model(ie_api_model);
             ov::pass::Manager manager;
             manager.register_pass<ov::pass::MakeStateful>(pairs_to_replace);
             manager.run_passes(model);
@@ -97,15 +104,17 @@ void regmodule_offline_transformations(py::module m) {
 
     m_offline_transformations.def(
         "compress_model_transformation",
-        [](std::shared_ptr<ov::Model> model) {
+        [](py::object& ie_api_model) {
+            const auto model = Common::utils::convert_to_model(ie_api_model);
             bool postponed = false;
             return ov::pass::compress_model_to_f16(model, postponed);
         },
         py::arg("model"));
 
     m_offline_transformations.def(
         "compress_quantize_weights_transformation",
-        [](std::shared_ptr<ov::Model> model) {
+        [](py::object& ie_api_model) {
+            const auto model = Common::utils::convert_to_model(ie_api_model);
             ov::pass::Manager manager;
             manager.register_pass<ov::pass::CompressQuantizeWeights>();
             manager.run_passes(model);
@@ -114,7 +123,8 @@ void regmodule_offline_transformations(py::module m) {
 
     m_offline_transformations.def(
         "convert_sequence_to_tensor_iterator_transformation",
-        [](std::shared_ptr<ov::Model> model) {
+        [](py::object ie_api_model) {
+            const auto model = Common::utils::convert_to_model(ie_api_model);
             ov::pass::Manager manager;
             manager.register_pass<ov::pass::ConvertSequenceToTensorIterator>();
             manager.run_passes(model);
@@ -123,7 +133,8 @@ void regmodule_offline_transformations(py::module m) {
 
     m_offline_transformations.def(
         "apply_fused_names_cleanup",
-        [](std::shared_ptr<ov::Model> model) {
+        [](py::object ie_api_model) {
+            const auto model = Common::utils::convert_to_model(ie_api_model);
             ov::pass::Manager manager;
             manager.register_pass<ov::pass::FusedNamesCleanup>();
             manager.run_passes(model);
@@ -132,7 +143,8 @@ void regmodule_offline_transformations(py::module m) {
 
     m_offline_transformations.def(
         "paged_attention_transformation",
-        [](std::shared_ptr<ov::Model> model, bool use_block_indices_inputs, bool use_score_outputs) {
+        [](py::object& ie_api_model, bool use_block_indices_inputs, bool use_score_outputs) {
+            const auto model = Common::utils::convert_to_model(ie_api_model);
             ov::pass::Manager manager;
             manager.register_pass<ov::pass::SDPAToPagedAttention>(use_block_indices_inputs, use_score_outputs);
             manager.run_passes(model);
@@ -143,7 +155,8 @@ void regmodule_offline_transformations(py::module m) {
 
     m_offline_transformations.def(
         "stateful_to_stateless_transformation",
-        [](std::shared_ptr<ov::Model> model) {
+        [](py::object& ie_api_model) {
+            const auto model = Common::utils::convert_to_model(ie_api_model);
             ov::pass::Manager manager;
             manager.register_pass<ov::pass::StatefulToStateless>();
             manager.run_passes(model);

src/bindings/python/src/pyopenvino/frontend/frontend.cpp

@@ -113,10 +113,13 @@ void regclass_frontend_FrontEnd(py::module m) {
                 :rtype: openvino.runtime.Model
              )");
 
-    fem.def("convert",
-            static_cast<void (FrontEnd::*)(const std::shared_ptr<ov::Model>&) const>(&FrontEnd::convert),
-            py::arg("model"),
-            R"(
+    fem.def(
+        "convert",
+        [](FrontEnd& self, const py::object& ie_api_model) {
+            return self.convert(Common::utils::convert_to_model(ie_api_model));
+        },
+        py::arg("model"),
+        R"(
                 Completely convert the remaining, not converted part of a function.
 
                 :param model: Partially converted OpenVINO model.
@@ -153,10 +156,13 @@ void regclass_frontend_FrontEnd(py::module m) {
                 :rtype: openvino.runtime.Model
              )");
 
-    fem.def("normalize",
-            &FrontEnd::normalize,
-            py::arg("model"),
-            R"(
+    fem.def(
+        "normalize",
+        [](FrontEnd& self, const py::object& ie_api_model) {
+            self.normalize(Common::utils::convert_to_model(ie_api_model));
+        },
+        py::arg("model"),
+        R"(
                 Runs normalization passes on function that was loaded with partial conversion.
 
                 :param model : Partially converted OpenVINO model.

src/bindings/python/src/pyopenvino/graph/attribute_visitor.cpp

@@ -35,6 +35,7 @@ void regclass_graph_AttributeVisitor(py::module m) {
         "on_attributes",
         [](ov::AttributeVisitor* self, py::dict& attributes) {
             py::object float_32_type = py::module_::import("numpy").attr("float32");
+            py::object model = py::module_::import("openvino.runtime").attr("Model");
             for (const auto& attribute : attributes) {
                 if (py::isinstance<py::bool_>(attribute.second)) {
                     visit_attribute<bool>(attributes, attribute, self);
@@ -48,6 +49,10 @@ void regclass_graph_AttributeVisitor(py::module m) {
                     visit_attribute<float>(attributes, attribute, self);
                 } else if (py::isinstance<ov::Model>(attribute.second)) {
                     visit_attribute<std::shared_ptr<ov::Model>>(attributes, attribute, self);
+                } else if (py::isinstance(attribute.second, model)) {
+                    auto attr_casted = attribute.second.attr("_Model__model").cast<std::shared_ptr<ov::Model>>();
+                    self->on_attribute<std::shared_ptr<ov::Model>>(attribute.first.cast<std::string>(), attr_casted);
+                    attributes[attribute.first] = std::move(attr_casted);
                 } else if (py::isinstance<ov::Dimension>(attribute.second)) {
                     visit_attribute<ov::Dimension>(attributes, attribute, self);
                 } else if (py::isinstance<ov::PartialShape>(attribute.second)) {

src/bindings/python/src/pyopenvino/graph/model.cpp

@@ -1328,10 +1328,6 @@ void regclass_graph_Model(py::module m) {
                outputs_str + "\n]>";
     });
 
-    model.def("__copy__", [](ov::Model& self) {
-        throw py::type_error("Cannot copy 'openvino.runtime.Model. Please, use deepcopy instead.");
-    });
-
     model.def("get_rt_info",
               (PyRTMap & (ov::Model::*)()) & ov::Model::get_rt_info,
               py::return_value_policy::reference_internal,

src/bindings/python/src/pyopenvino/graph/ops/if.cpp

@@ -12,6 +12,7 @@
 #include "pyopenvino/core/common.hpp"
 #include "pyopenvino/graph/ops/if.hpp"
 #include "pyopenvino/graph/ops/util/multisubgraph.hpp"
+#include "pyopenvino/utils/utils.hpp"
 
 namespace py = pybind11;
 
@@ -77,10 +78,14 @@ void regclass_graph_op_If(py::module m) {
             :rtype: openvino.Model
         )");
 
-    cls.def("set_then_body",
-            &ov::op::v8::If::set_then_body,
-            py::arg("body"),
-            R"(
+    cls.def(
+        "set_then_body",
+        [](const std::shared_ptr<ov::op::v8::If>& self, const py::object& ie_api_model) {
+            const auto body = Common::utils::convert_to_model(ie_api_model);
+            return self->set_then_body(body);
+        },
+        py::arg("body"),
+        R"(
             Sets new Model object as new then_body.
 
             :param body: new body for 'then' branch.
@@ -89,10 +94,14 @@ void regclass_graph_op_If(py::module m) {
             :rtype: None
         )");
 
-    cls.def("set_else_body",
-            &ov::op::v8::If::set_else_body,
-            py::arg("body"),
-            R"(
+    cls.def(
+        "set_else_body",
+        [](const std::shared_ptr<ov::op::v8::If>& self, const py::object& ie_api_model) {
+            const auto body = Common::utils::convert_to_model(ie_api_model);
+            return self->set_else_body(body);
+        },
+        py::arg("body"),
+        R"(
             Sets new Model object as new else_body.
 
             :param body: new body for 'else' branch.
@@ -156,11 +165,15 @@ void regclass_graph_op_If(py::module m) {
             :rtype: openvino.Model
         )");
 
-    cls.def("set_function",
-            &ov::op::util::MultiSubGraphOp::set_function,
-            py::arg("index"),
-            py::arg("func"),
-            R"(
+    cls.def(
+        "set_function",
+        [](const std::shared_ptr<ov::op::v8::If>& self, int index, const py::object& ie_api_model) {
+            const auto func = Common::utils::convert_to_model(ie_api_model);
+            self->set_function(index, func);
+        },
+        py::arg("index"),
+        py::arg("func"),
+        R"(
             Adds sub-graph to MultiSubGraphOp.
 
             :param index: index of new sub-graph.