diff --git a/__init__.py b/__init__.py
index 7a01ac8..293366f 100644
--- a/__init__.py
+++ b/__init__.py
@@ -1,6 +1,4 @@
 """Top-level package for basic_data_handling."""
-import os, sys
-sys.path.append(os.path.dirname(os.path.realpath(__file__)))
 
 __all__ = [
     "NODE_CLASS_MAPPINGS",
@@ -8,7 +6,14 @@
     "WEB_DIRECTORY",
 ]
 
-from src.basic_data_handling import NODE_CLASS_MAPPINGS
-from src.basic_data_handling import NODE_DISPLAY_NAME_MAPPINGS
+try:
+    # For ComfyUI
+    from .src.basic_data_handling import NODE_CLASS_MAPPINGS
+    from .src.basic_data_handling import NODE_DISPLAY_NAME_MAPPINGS
+except ImportError:
+    # For running tests
+    from src.basic_data_handling import NODE_CLASS_MAPPINGS
+    from src.basic_data_handling import NODE_DISPLAY_NAME_MAPPINGS
+
 
 WEB_DIRECTORY = "./web"
diff --git a/pyproject.toml b/pyproject.toml
index 2b1e18c..b81f529 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "basic_data_handling"
-version = "0.5.0"
+version = "0.6.0"
 description = """Basic Python functions for manipulating data that every programmer is used to, lightweight with no additional dependencies.
 
 Supported data types: