diff --git a/comfy/model_management.py b/comfy/model_management.py
index 143991ad6..70c2d5e22 100644
--- a/comfy/model_management.py
+++ b/comfy/model_management.py
@@ -27,6 +27,7 @@ import weakref
 import gc
 import os
 from contextlib import nullcontext
+import comfy.memory_management
 import comfy.utils
 import comfy.quant_ops