load python package from mecha_extensions

__init__.py

@@ -3,17 +3,37 @@ import importlib.util
 import sys
 import traceback
 
+
 extensions_path = pathlib.Path(__file__).parent / "mecha_extensions"
 for path in extensions_path.iterdir():
-    if path.is_file() and path.suffix == ".py":
-        try:
-            module_spec = importlib.util.spec_from_file_location(path.stem, str(path))
-            module = importlib.util.module_from_spec(module_spec)
-            module_spec.loader.exec_module(module)
-        except Exception as e:
-            print(f"[comfy-mecha] could not import mecha extension {path}: {e}", file=sys.stderr)
-            traceback.print_exc()
+    is_script = path.is_file() and path.suffix == ".py"
+    is_package = path.is_dir() and (path / "__init__.py").is_file()
+
+    if not (is_script or is_package):
+        continue
+
+    if is_package:
+        script_path = path / "__init__.py"
+        module_name = f"mecha_extensions.{path.name}"
+    else:
+        script_path = path
+        module_name = f"mecha_extensions.{path.stem}"
+
+    try:
+        spec = importlib.util.spec_from_file_location(module_name, str(script_path))
+        if spec is None or spec.loader is None:
+            raise ImportError(f"Could not create a module spec for {script_path!r}")
+
+        module = importlib.util.module_from_spec(spec)
+        sys.modules[module_name] = module
+        spec.loader.exec_module(module)
+    except Exception as e:
+        print(f"[comfy-mecha] could not import mecha extension {path}: {e}", file=sys.stderr)
+        traceback.print_exc()
+    finally:
+        if module_name in sys.modules:
+            del sys.modules[module_name]
+
 
 from .nodes import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS
-
 WEB_DIRECTORY = "js"

nodes/merge_nodes.py

@@ -582,7 +582,7 @@ def make_comfy_node_class(class_name: str, method: MergeMethod) -> type:
                                "you probably want to either use this *or* a cache unit, but not both at the same time "
                                "because the memory adds up.\n\n"
                                "The difference between merge checkpointing and cache is that merge checkpointing "
-                               "completely re-merges from scratch if any input changes."
+                               "completely re-merges from scratch if any input changes. "
                                "Merge checkpointing is also generally much faster than cache in the fast path.",
                 }),
             },

pyproject.toml

@@ -1,7 +1,7 @@
 [project]
 name = "comfy-mecha"
 description = "model merging nodes powered by sd-mecha, a memory efficient merging library"
-version = "1.2.35"
+version = "1.2.36"
 license = { file = "LICENSE" }
 dependencies = ["sd-mecha==1.0.4"]