diff --git a/SampleDiffusion/EXT_SampleDiffusion.py b/SampleDiffusion/EXT_SampleDiffusion.py
index d41f1b6..fb2bfc5 100644
--- a/SampleDiffusion/EXT_SampleDiffusion.py
+++ b/SampleDiffusion/EXT_SampleDiffusion.py
@@ -8,6 +8,7 @@
 from comfy.model_management import get_torch_device
 import importlib
 import yaml
+from pathlib import Path
 
 def hijack_import(importname, installname):
     try:
@@ -107,7 +108,7 @@ async def view_image(request):
 # init and sample_diffusion lib load
 
 
-comfy_dir = get_comfy_dir()   
+comfy_dir = Path(__file__).resolve().parent.parent
 if not os.path.exists(os.path.join(comfy_dir, 'custom_nodes/SampleDiffusion/libs')):
     os.makedirs(os.path.join(comfy_dir, 'custom_nodes/SampleDiffusion/libs'))
 lib = os.path.join(comfy_dir, 'custom_nodes/SampleDiffusion/libs/sample_generator') 
diff --git a/SampleDiffusion/EXT_VariationUtils.py b/SampleDiffusion/EXT_VariationUtils.py
index 530c5f3..570261d 100644
--- a/SampleDiffusion/EXT_VariationUtils.py
+++ b/SampleDiffusion/EXT_VariationUtils.py
@@ -4,7 +4,7 @@
 import soundfile as sf
 
 from comfy.model_management import get_torch_device
-from custom_nodes.SampleDiffusion.EXT_SampleDiffusion import AudioInference
+from .EXT_SampleDiffusion import AudioInference
 from diffusion_library.sampler import SamplerType
 from diffusion_library.scheduler import SchedulerType
 
diff --git a/SampleDiffusion/__init__.py b/SampleDiffusion/__init__.py
index 80d6c15..13cb64a 100644
--- a/SampleDiffusion/__init__.py
+++ b/SampleDiffusion/__init__.py
@@ -3,9 +3,11 @@
 
 NODE_CLASS_MAPPINGS = {}
 
+current_package = __name__
+
 for node in os.listdir(os.path.dirname(__file__)):
     if node.startswith('EXT_'):
         node = node.split('.')[0]
-        node_import = importlib.import_module('custom_nodes.SampleDiffusion.' + node)
+        node_import = importlib.import_module(f"{current_package}.{node}")
         # get class node mappings from py file
         NODE_CLASS_MAPPINGS.update(node_import.NODE_CLASS_MAPPINGS)
diff --git a/__init__.py b/__init__.py
new file mode 100644
index 0000000..95f0eb4
--- /dev/null
+++ b/__init__.py
@@ -0,0 +1 @@
+from .SampleDiffusion import *