From 33e71e0e7925442bf1c261c8d921f2c541606529 Mon Sep 17 00:00:00 2001 From: ethan Date: Fri, 24 Jan 2025 01:37:44 -0800 Subject: [PATCH 1/7] update openvino backend --- comfy/cli_args.py | 1 + comfy/model_management.py | 2 +- comfy/model_patcher.py | 9 +++++++++ comfy/sampler_helpers.py | 7 ++++++- 4 files changed, 17 insertions(+), 2 deletions(-) diff --git a/comfy/cli_args.py b/comfy/cli_args.py index 812798bf8..369265818 100644 --- a/comfy/cli_args.py +++ b/comfy/cli_args.py @@ -83,6 +83,7 @@ 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/model_management.py b/comfy/model_management.py index f6dfc18b0..29aa60d68 100644 --- a/comfy/model_management.py +++ b/comfy/model_management.py @@ -93,7 +93,7 @@ try: except: npu_available = False -if args.cpu: +if args.cpu or args.openvino: cpu_state = CPUState.CPU def is_intel_xpu(): diff --git a/comfy/model_patcher.py b/comfy/model_patcher.py index 0501f7b38..813e7be12 100644 --- a/comfy/model_patcher.py +++ b/comfy/model_patcher.py @@ -469,6 +469,15 @@ class ModelPatcher: current_patches = self.patches.get(key, []) current_patches.append((strength_patch, patches[k], strength_model, offset, function)) self.patches[key] = current_patches + else: + new_key=key.replace("diffusion_model","diffusion_model._orig_mod") + if new_key in model_sd: + p.add(k) + if key in self.patches: + self.patches.pop(key) + current_patches = self.patches.get(new_key, []) + current_patches.append((strength_patch, patches[k], strength_model, offset, function)) + self.patches[new_key] = current_patches self.patches_uuid = uuid.uuid4() return list(p) diff --git a/comfy/sampler_helpers.py b/comfy/sampler_helpers.py index 92ec7ca7a..dc7e07aaa 100644 --- a/comfy/sampler_helpers.py +++ b/comfy/sampler_helpers.py @@ -1,4 +1,5 @@ from __future__ import annotations +from comfy.cli_args import args import uuid import comfy.model_management import comfy.conds @@ -114,7 +115,11 @@ 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): From d1f61cca5e4a3761ee6897655fa4d422ae9b8b47 Mon Sep 17 00:00:00 2001 From: ethan Date: Wed, 29 Jan 2025 07:03:31 -0800 Subject: [PATCH 2/7] 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 369265818..812798bf8 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 dc7e07aaa..8c080baf6 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 1fe6f42c7..611e4f9e2 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, From 317af7201f54e464bc35c5dc5b6c42761df2051c Mon Sep 17 00:00:00 2001 From: ethan Date: Wed, 29 Jan 2025 07:05:47 -0800 Subject: [PATCH 3/7] remove history commit remove history commit remove history commit --- comfy/model_management.py | 2 +- comfy/model_patcher.py | 2 +- comfy/sampler_helpers.py | 2 +- 3 files changed, 3 insertions(+), 3 deletions(-) diff --git a/comfy/model_management.py b/comfy/model_management.py index 29aa60d68..f6dfc18b0 100644 --- a/comfy/model_management.py +++ b/comfy/model_management.py @@ -93,7 +93,7 @@ try: except: npu_available = False -if args.cpu or args.openvino: +if args.cpu: cpu_state = CPUState.CPU def is_intel_xpu(): diff --git a/comfy/model_patcher.py b/comfy/model_patcher.py index 813e7be12..a2a26f174 100644 --- a/comfy/model_patcher.py +++ b/comfy/model_patcher.py @@ -470,7 +470,7 @@ class ModelPatcher: current_patches.append((strength_patch, patches[k], strength_model, offset, function)) self.patches[key] = current_patches else: - new_key=key.replace("diffusion_model","diffusion_model._orig_mod") + new_key=key.replace("diffusion_model", "diffusion_model._orig_mod") if new_key in model_sd: p.add(k) if key in self.patches: diff --git a/comfy/sampler_helpers.py b/comfy/sampler_helpers.py index 8c080baf6..92ec7ca7a 100644 --- a/comfy/sampler_helpers.py +++ b/comfy/sampler_helpers.py @@ -1,5 +1,4 @@ from __future__ import annotations -from comfy.cli_args import args import uuid import comfy.model_management import comfy.conds @@ -115,6 +114,7 @@ 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 + return real_model, conds, models def cleanup_models(conds, models): From 77e9294c08351acabb1efdc3c10b021770f71807 Mon Sep 17 00:00:00 2001 From: ethan Date: Thu, 30 Jan 2025 00:20:58 -0800 Subject: [PATCH 4/7] add Query Device --- comfy_extras/nodes_torch_compile.py | 13 +++++++++++-- 1 file changed, 11 insertions(+), 2 deletions(-) diff --git a/comfy_extras/nodes_torch_compile.py b/comfy_extras/nodes_torch_compile.py index 611e4f9e2..cad3b1422 100644 --- a/comfy_extras/nodes_torch_compile.py +++ b/comfy_extras/nodes_torch_compile.py @@ -1,16 +1,25 @@ import torch +import importlib class TorchCompileModel: @classmethod def INPUT_TYPES(s): + if importlib.util.find_spec("openvino") is not None: + import openvino as ov + + core = ov.Core() + available_devices = core.available_devices + else: + available_devices = [] + return { "required": { "model": ("MODEL",), "backend": (["inductor", "cudagraphs", "openvino"],), }, "optional": { - "openvino_device": (["CPU", "GPU", "NPU"],), + "openvino_device": (available_devices,), }, } @@ -46,4 +55,4 @@ class TorchCompileModel: NODE_CLASS_MAPPINGS = { "TorchCompileModel": TorchCompileModel, -} +} \ No newline at end of file From ebfe7a5679e5200527fc18eb1f56115ca0ddd3f4 Mon Sep 17 00:00:00 2001 From: ethan Date: Mon, 10 Feb 2025 19:54:44 -0800 Subject: [PATCH 5/7] fix the issue for model first inference with lora --- comfy/model_patcher.py | 14 +++----------- 1 file changed, 3 insertions(+), 11 deletions(-) diff --git a/comfy/model_patcher.py b/comfy/model_patcher.py index a2a26f174..31c98a041 100644 --- a/comfy/model_patcher.py +++ b/comfy/model_patcher.py @@ -463,21 +463,13 @@ class ModelPatcher: key = k[0] if len(k) > 2: function = k[2] - - if key in model_sd: + org_key=key.replace("diffusion_model", "diffusion_model._orig_mod") + if key in model_sd or org_key in model_sd: p.add(k) current_patches = self.patches.get(key, []) current_patches.append((strength_patch, patches[k], strength_model, offset, function)) self.patches[key] = current_patches - else: - new_key=key.replace("diffusion_model", "diffusion_model._orig_mod") - if new_key in model_sd: - p.add(k) - if key in self.patches: - self.patches.pop(key) - current_patches = self.patches.get(new_key, []) - current_patches.append((strength_patch, patches[k], strength_model, offset, function)) - self.patches[new_key] = current_patches + self.patches[org_key] = current_patches self.patches_uuid = uuid.uuid4() return list(p) From bc9eb9dfdb571f0ca0248691f1f5850e32ec3221 Mon Sep 17 00:00:00 2001 From: ethan Date: Wed, 5 Mar 2025 18:28:47 -0800 Subject: [PATCH 6/7] fix the memory leakage issue --- comfy/model_patcher.py | 9 ++++++++- comfy_extras/nodes_torch_compile.py | 11 +++++++++-- 2 files changed, 17 insertions(+), 3 deletions(-) diff --git a/comfy/model_patcher.py b/comfy/model_patcher.py index 259d66d9d..9a8f7fe78 100644 --- a/comfy/model_patcher.py +++ b/comfy/model_patcher.py @@ -499,12 +499,19 @@ class ModelPatcher: if len(k) > 2: function = k[2] org_key=key.replace("diffusion_model", "diffusion_model._orig_mod") - if key in model_sd or org_key in model_sd: + if key in model_sd: p.add(k) current_patches = self.patches.get(key, []) current_patches.append((strength_patch, patches[k], strength_model, offset, function)) self.patches[key] = current_patches self.patches[org_key] = current_patches + elif org_key in model_sd: + if key in self.patches: + self.patches.pop(key) + p.add(k) + current_patches = self.patches.get(org_key, []) + current_patches.append((strength_patch, patches[k], strength_model, offset, function)) + self.patches[org_key] = current_patches self.patches_uuid = uuid.uuid4() return list(p) diff --git a/comfy_extras/nodes_torch_compile.py b/comfy_extras/nodes_torch_compile.py index cad3b1422..d217f5041 100644 --- a/comfy_extras/nodes_torch_compile.py +++ b/comfy_extras/nodes_torch_compile.py @@ -19,7 +19,7 @@ class TorchCompileModel: "backend": (["inductor", "cudagraphs", "openvino"],), }, "optional": { - "openvino_device": (available_devices,), + "openvino device": (available_devices,), }, } @@ -30,6 +30,7 @@ class TorchCompileModel: EXPERIMENTAL = True def patch(self, model, backend, openvino_device): + print(model.__class__.__name__) if backend == "openvino": options = {"device": openvino_device} try: @@ -39,6 +40,12 @@ class TorchCompileModel: "Could not import openvino python package. " "Please install it with `pip install openvino`." ) + import openvino.frontend.pytorch.torchdynamo.execute as ov_ex + + torch._dynamo.reset() + ov_ex.compiled_cache.clear() + ov_ex.req_cache.clear() + ov_ex.partitioned_modules.clear() else: options = None m = model.clone() @@ -55,4 +62,4 @@ class TorchCompileModel: NODE_CLASS_MAPPINGS = { "TorchCompileModel": TorchCompileModel, -} \ No newline at end of file +} From e14b8dfec5b49f691aedb1e5139e5d8d9015bbef Mon Sep 17 00:00:00 2001 From: Ethan Yang Date: Fri, 7 Mar 2025 13:23:19 +0800 Subject: [PATCH 7/7] Update nodes_torch_compile.py --- comfy_extras/nodes_torch_compile.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/comfy_extras/nodes_torch_compile.py b/comfy_extras/nodes_torch_compile.py index d217f5041..8acfe4c11 100644 --- a/comfy_extras/nodes_torch_compile.py +++ b/comfy_extras/nodes_torch_compile.py @@ -19,7 +19,7 @@ class TorchCompileModel: "backend": (["inductor", "cudagraphs", "openvino"],), }, "optional": { - "openvino device": (available_devices,), + "openvino_device": (available_devices,), }, }