From 6c59d5c32d7e672de9dffb8727d42e96f25a4685 Mon Sep 17 00:00:00 2001
From: Mel Massadian <mel@melmassadian.com>
Date: Tue, 24 Jun 2025 17:08:03 +0200
Subject: [PATCH] =?UTF-8?q?chore:=20=F0=9F=A7=B9=20support=20hot=20reloadi?=
 =?UTF-8?q?ng?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 .gitignore  |  1 +
 __init__.py | 34 ++++++++++++++++++++++++++--------
 2 files changed, 27 insertions(+), 8 deletions(-)

diff --git a/.gitignore b/.gitignore
index 26ea882..6a48272 100644
--- a/.gitignore
+++ b/.gitignore
@@ -13,3 +13,4 @@ node.zip
 
 # I store the gh-pages worktrees (src & build) there
 .worktrees
+comfy.lock
diff --git a/__init__.py b/__init__.py
index c973abf..e759cea 100644
--- a/__init__.py
+++ b/__init__.py
@@ -34,6 +34,8 @@ from aiohttp import web
 
 IN_COMFY = False
 
+PromptServer = None
+
 try:
     from server import PromptServer
 
@@ -240,10 +242,29 @@ if failed:
 # - ENDPOINT
 
 
-if IN_COMFY and hasattr(PromptServer, "instance"):
+# TODO: move that away and simplify existing endpoints
+
+
+def register_routes():
+    if not PromptServer:
+        log.error("No prompt server, are you inside comfy?")
+
+    if PromptServer.instance.app.frozen:
+        log.warning(
+            "The router is frozen and cannot be further edited."
+            "If you are hot reloading mtb this is expected."
+        )
+        return
+
     img_cache = None
     prompt_cache = None
 
+    import asyncio
+    import os
+    from io import BytesIO
+
+    from PIL import Image
+
     with contextlib.suppress(ImportError):
         from cachetools import TTLCache
 
@@ -360,13 +381,6 @@ if IN_COMFY and hasattr(PromptServer, "instance"):
         # Return JSON for other requests
         return web.json_response({"message": "Welcome to MTB!"})
 
-    import asyncio
-    import os
-    from io import BytesIO
-
-    from aiohttp import web
-    from PIL import Image
-
     def get_cached_image(file_path: str, preview_params=None, channel=None):
         cache_key = (file_path, preview_params, channel)
         if img_cache and (cache_key in img_cache):
@@ -571,6 +585,10 @@ if IN_COMFY and hasattr(PromptServer, "instance"):
         return await endpoint.do_action(request)
 
 
+if IN_COMFY and hasattr(PromptServer, "instance"):
+    register_routes()
+
+
 # - WAS Dictionary
 MANIFEST = {
     "name": "MTB Nodes",  # The title that will be displayed on Node Class menu,. and Node Class view