diff --git a/nodes/image_nodes.py b/nodes/image_nodes.py
index f910251..2a2879c 100644
--- a/nodes/image_nodes.py
+++ b/nodes/image_nodes.py
@@ -23,7 +23,7 @@ from comfy_extras.nodes_mask import ImageCompositeMasked
 from comfy.cli_args import args
 from comfy.utils import ProgressBar, common_upscale
 import folder_paths
-import model_management
+from comfy import model_management
 
 script_directory = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
 
diff --git a/nodes/mask_nodes.py b/nodes/mask_nodes.py
index ba92228..ce406cc 100644
--- a/nodes/mask_nodes.py
+++ b/nodes/mask_nodes.py
@@ -7,7 +7,7 @@ import numpy as np
 from contextlib import nullcontext
 import os
 
-import model_management
+from comfy import model_management
 from comfy.utils import ProgressBar
 from comfy.utils import common_upscale
 from nodes import MAX_RESOLUTION
diff --git a/nodes/nodes.py b/nodes/nodes.py
index 05b21e4..38d6d07 100644
--- a/nodes/nodes.py
+++ b/nodes/nodes.py
@@ -7,7 +7,7 @@ import json, re, os, io, time, platform
 import re
 import importlib
 
-import model_management
+from comfy import model_management
 import folder_paths
 from nodes import MAX_RESOLUTION
 from comfy.utils import common_upscale, ProgressBar, load_torch_file