diff --git a/comfy/cmd/server.py b/comfy/cmd/server.py
index 7eb511edb..d93bc63f5 100644
--- a/comfy/cmd/server.py
+++ b/comfy/cmd/server.py
@@ -68,7 +68,7 @@ class HeuristicPath(NamedTuple):
 
 
 # Import cache control middleware
-from middleware.cache_middleware import cache_control
+from ..middleware.cache_middleware import cache_control
 
 async def send_socket_catch_exception(function, message):
     try:
diff --git a/middleware/__init__.py b/comfy/middleware/__init__.py
similarity index 100%
rename from middleware/__init__.py
rename to comfy/middleware/__init__.py
diff --git a/middleware/cache_middleware.py b/comfy/middleware/cache_middleware.py
similarity index 100%
rename from middleware/cache_middleware.py
rename to comfy/middleware/cache_middleware.py