From 257088376980c477cba763996255b42310de6eb2 Mon Sep 17 00:00:00 2001
From: Helena <helena.kloosterman@intel.com>
Date: Fri, 17 Jan 2025 09:21:32 +0100
Subject: [PATCH 1/2] openvino.runtime -> openvino in import

---
 tools/llm_bench/benchmark.py                        | 2 +-
 tools/llm_bench/convert.py                          | 2 +-
 tools/llm_bench/llm_bench_utils/ov_model_classes.py | 2 +-
 tools/llm_bench/llm_bench_utils/ov_utils.py         | 4 ++--
 4 files changed, 5 insertions(+), 5 deletions(-)

diff --git a/tools/llm_bench/benchmark.py b/tools/llm_bench/benchmark.py
index b9e038c12e..6fc135c4ef 100644
--- a/tools/llm_bench/benchmark.py
+++ b/tools/llm_bench/benchmark.py
@@ -6,7 +6,7 @@
 import argparse
 import logging as log
 import llm_bench_utils.model_utils
-from openvino.runtime import get_version
+from openvino import get_version
 import torch
 import traceback
 from llm_bench_utils.memory_profile import MemConsumption
diff --git a/tools/llm_bench/convert.py b/tools/llm_bench/convert.py
index aa9ab504de..732e9407c8 100644
--- a/tools/llm_bench/convert.py
+++ b/tools/llm_bench/convert.py
@@ -20,7 +20,7 @@
 from diffusers import UNet2DConditionModel, AutoencoderTiny, LCMScheduler
 from nncf.torch.model_creation import is_wrapped_model
 from openvino import Type as OVType, PartialShape, save_model, convert_model
-from openvino.runtime import Core, get_version
+from openvino import Core, get_version
 from optimum.exporters import TasksManager
 from optimum.utils import DEFAULT_DUMMY_SHAPES
 from optimum.intel.openvino.configuration import OVConfig
diff --git a/tools/llm_bench/llm_bench_utils/ov_model_classes.py b/tools/llm_bench/llm_bench_utils/ov_model_classes.py
index 17cfab4e02..ccb4e9af1f 100644
--- a/tools/llm_bench/llm_bench_utils/ov_model_classes.py
+++ b/tools/llm_bench/llm_bench_utils/ov_model_classes.py
@@ -16,7 +16,7 @@
 from diffusers.pipelines.pipeline_utils import DiffusionPipeline, ImagePipelineOutput
 from optimum.intel.openvino import OVModelForCausalLM
 from optimum.intel.openvino.utils import ONNX_WEIGHTS_NAME, OV_XML_FILE_NAME
-from openvino.runtime import Model, Core, Tensor, Type
+from openvino import Model, Core, Tensor, Type
 from transformers import PretrainedConfig
 from transformers.modeling_outputs import CausalLMOutputWithPast, ModelOutput
 from transformers import GenerationConfig, StoppingCriteriaList
diff --git a/tools/llm_bench/llm_bench_utils/ov_utils.py b/tools/llm_bench/llm_bench_utils/ov_utils.py
index 606a39cecf..9a84388e9c 100644
--- a/tools/llm_bench/llm_bench_utils/ov_utils.py
+++ b/tools/llm_bench/llm_bench_utils/ov_utils.py
@@ -3,7 +3,7 @@
 # SPDX-License-Identifier: Apache-2.0
 from pathlib import Path
 from transformers import AutoConfig, AutoProcessor, AutoTokenizer
-from openvino.runtime import Core
+from openvino import Core
 import openvino as ov
 import logging as log
 import torch
@@ -17,7 +17,7 @@
     DEFAULT_MODEL_CLASSES,
     IMAGE_GEN_CLS
 )
-import openvino.runtime.opset13 as opset
+import openvino.opset13 as opset
 from transformers import pipeline
 import openvino_genai as ov_genai
 import queue

From 49cc3d56b2773f707da782ed255f6ee998616c0e Mon Sep 17 00:00:00 2001
From: Helena <helena.kloosterman@intel.com>
Date: Fri, 17 Jan 2025 13:08:48 +0100
Subject: [PATCH 2/2] Add try/except around openvino.opset13 import

---
 tools/llm_bench/llm_bench_utils/ov_utils.py | 5 ++++-
 1 file changed, 4 insertions(+), 1 deletion(-)

diff --git a/tools/llm_bench/llm_bench_utils/ov_utils.py b/tools/llm_bench/llm_bench_utils/ov_utils.py
index 9a84388e9c..c70e4beb5e 100644
--- a/tools/llm_bench/llm_bench_utils/ov_utils.py
+++ b/tools/llm_bench/llm_bench_utils/ov_utils.py
@@ -17,7 +17,10 @@
     DEFAULT_MODEL_CLASSES,
     IMAGE_GEN_CLS
 )
-import openvino.opset13 as opset
+try:
+    import openvino.opset13 as opset
+except ImportError:
+    import openvino.runtime.opset13 as opset
 from transformers import pipeline
 import openvino_genai as ov_genai
 import queue