diff --git a/__init__.py b/__init__.py
index dda039b..cf0f494 100644
--- a/__init__.py
+++ b/__init__.py
@@ -1,5 +1,5 @@
 import os
-import model_management
+from comfy import model_management
 import torch
 import comfy.utils
 import numpy as np