add model caching

Author: openvino-dev-samples

node_openvino.py

@@ -1,7 +1,7 @@
 import torch
 import openvino as ov
 import openvino.frontend.pytorch.torchdynamo.execute as ov_ex
-
+import os
 from comfy_api.torch_helpers import set_torch_compile_wrapper
 
 
@@ -10,11 +10,12 @@ class TorchCompileModelOpenVINO:
     def INPUT_TYPES(s):
         core = ov.Core()
         available_devices = core.available_devices
-
+        model_cache_option = ["OFF", "ON"]
         return {
             "required": {
                 "model": ("MODEL",),
                 "device": (available_devices,),
+                "model_cache": (model_cache_option,),
             },
         }
 
@@ -24,21 +25,32 @@ def INPUT_TYPES(s):
     CATEGORY = "OpenVINO"
     EXPERIMENTAL = True
 
-    def patch(self, model, device):
-        options = {"device": device}
+    def patch(self, model, device, model_cache):
+        model_cache_option = {"OFF": False, "ON": True}
+        cache_path = os.path.join(os.getcwd(), "openvino_model_cache")
+        options = {
+            "device": device,
+            "model_caching": model_cache_option[model_cache],
+            "cache_dir": cache_path,
+        }
         torch._dynamo.reset()
         ov_ex.compiled_cache.clear()
         ov_ex.req_cache.clear()
         ov_ex.partitioned_modules.clear()
         m = model.clone()
-        set_torch_compile_wrapper(m,
-                                  backend="openvino",
-                                  options=options,
-                                  )
+
+        if model_cache_option[model_cache]:
+            print(f"Using the model cache from {cache_path}.")
+
+        set_torch_compile_wrapper(
+            m,
+            backend="openvino",
+            options=options,
+        )
         return (m,)
 
 
 # The node ID in NODE_CLASS_MAPPINGS should be globally unique across ComfyUI ecosystem
 NODE_CLASS_MAPPINGS = {
     "OpenVINO_TorchCompileModel": TorchCompileModelOpenVINO,
-}
+}

pyproject.toml

@@ -1,7 +1,7 @@
 [project]
 name = "comfyui-openvino"
 description = "OpenVINO node is designed for optimizing the performance of model inference in ComfyUI by leveraging Intel OpenVINO toolkits. It can support running model on Intel CPU, GPU and NPU device."
-version = "1.0.2"
+version = "1.0.3"
 license = {file = "LICENSE"}
 
 [project.urls]