diff --git a/comfy/model_patcher.py b/comfy/model_patcher.py index b7cb12dfc..2f76276fd 100644 --- a/comfy/model_patcher.py +++ b/comfy/model_patcher.py @@ -498,12 +498,20 @@ class ModelPatcher: key = k[0] if len(k) > 2: function = k[2] - + org_key=key.replace("diffusion_model", "diffusion_model._orig_mod") 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 1fe6f42c7..8acfe4c11 100644 --- a/comfy_extras/nodes_torch_compile.py +++ b/comfy_extras/nodes_torch_compile.py @@ -1,21 +1,64 @@ import torch +import importlib + class TorchCompileModel: @classmethod def INPUT_TYPES(s): - return {"required": { "model": ("MODEL",), - "backend": (["inductor", "cudagraphs"],), - }} + 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": (available_devices,), + }, + } + RETURN_TYPES = ("MODEL",) FUNCTION = "patch" CATEGORY = "_for_testing" EXPERIMENTAL = True - def patch(self, model, backend): + def patch(self, model, backend, openvino_device): + print(model.__class__.__name__) + 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`." + ) + 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() - 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,