From d1f61cca5e4a3761ee6897655fa4d422ae9b8b47 Mon Sep 17 00:00:00 2001
From: ethan <ethan.yang@intel.com>
Date: Wed, 29 Jan 2025 07:03:31 -0800
Subject: [PATCH] add openvino to torch compile

---
 comfy/cli_args.py                   |  1 -
 comfy/sampler_helpers.py            |  5 ----
 comfy_extras/nodes_torch_compile.py | 39 ++++++++++++++++++++++++-----
 3 files changed, 33 insertions(+), 12 deletions(-)

diff --git a/comfy/cli_args.py b/comfy/cli_args.py
index 36926581..812798bf 100644
--- a/comfy/cli_args.py
+++ b/comfy/cli_args.py
@@ -83,7 +83,6 @@ fpte_group.add_argument("--fp32-text-enc", action="store_true", help="Store text
 parser.add_argument("--force-channels-last", action="store_true", help="Force channels last format when inferencing the models.")
 
 parser.add_argument("--directml", type=int, nargs="?", metavar="DIRECTML_DEVICE", const=-1, help="Use torch-directml.")
-parser.add_argument("--openvino", type=str, default="GPU", help="Run OpenVINO inference engine on the specified device.")
 
 parser.add_argument("--oneapi-device-selector", type=str, default=None, metavar="SELECTOR_STRING", help="Sets the oneAPI device(s) this instance will use.")
 parser.add_argument("--disable-ipex-optimize", action="store_true", help="Disables ipex.optimize default when loading models with Intel's Extension for Pytorch.")
diff --git a/comfy/sampler_helpers.py b/comfy/sampler_helpers.py
index dc7e07aa..8c080baf 100644
--- a/comfy/sampler_helpers.py
+++ b/comfy/sampler_helpers.py
@@ -115,11 +115,6 @@ def prepare_sampling(model: ModelPatcher, noise_shape, conds, model_options=None
     minimum_memory_required = model.memory_required([noise_shape[0]] + list(noise_shape[1:])) + inference_memory
     comfy.model_management.load_models_gpu([model] + models, memory_required=memory_required, minimum_memory_required=minimum_memory_required)
     real_model = model.model
-    if args.openvino and real_model.diffusion_model.__class__.__name__=="UNetModel":
-        import openvino.torch
-        import torch
-        print("Unet is being compiled using OpenVINO")
-        real_model.diffusion_model = torch.compile(real_model.diffusion_model, backend="openvino", options = {"device" : args.openvino, "model_caching" : False, "cache_dir": "./model_cache"})
     return real_model, conds, models
 
 def cleanup_models(conds, models):
diff --git a/comfy_extras/nodes_torch_compile.py b/comfy_extras/nodes_torch_compile.py
index 1fe6f42c..611e4f9e 100644
--- a/comfy_extras/nodes_torch_compile.py
+++ b/comfy_extras/nodes_torch_compile.py
@@ -1,21 +1,48 @@
 import torch
 
+
 class TorchCompileModel:
     @classmethod
     def INPUT_TYPES(s):
-        return {"required": { "model": ("MODEL",),
-                             "backend": (["inductor", "cudagraphs"],),
-                              }}
+        return {
+            "required": {
+                "model": ("MODEL",),
+                "backend": (["inductor", "cudagraphs", "openvino"],),
+            },
+            "optional": {
+                "openvino_device": (["CPU", "GPU", "NPU"],),
+            },
+        }
+
     RETURN_TYPES = ("MODEL",)
     FUNCTION = "patch"
 
     CATEGORY = "_for_testing"
     EXPERIMENTAL = True
 
-    def patch(self, model, backend):
+    def patch(self, model, backend, openvino_device):
+        if backend == "openvino":
+            options = {"device": openvino_device}
+            try:
+                import openvino.torch
+            except ImportError:
+                raise ImportError(
+                    "Could not import openvino python package. "
+                    "Please install it with `pip install openvino`."
+                )
+        else:
+            options = None
         m = model.clone()
-        m.add_object_patch("diffusion_model", torch.compile(model=m.get_model_object("diffusion_model"), backend=backend))
-        return (m, )
+        m.add_object_patch(
+            "diffusion_model",
+            torch.compile(
+                model=m.get_model_object("diffusion_model"),
+                backend=backend,
+                options=options,
+            ),
+        )
+        return (m,)
+
 
 NODE_CLASS_MAPPINGS = {
     "TorchCompileModel": TorchCompileModel,