diff --git a/depthfm/__init__.py b/depthfm/__init__.py
index f2a14c2..b5853c0 100644
--- a/depthfm/__init__.py
+++ b/depthfm/__init__.py
@@ -1,5 +1,5 @@
-import os
-import sys
-sys.path.append(os.path.dirname(os.path.abspath(__file__)))
-from dfm import DepthFM
-from unet import UNetModel
+#import os
+#import sys
+#sys.path.append(os.path.dirname(os.path.abspath(__file__)))
+#from .dfm import DepthFM
+#from .unet import UNetModel
diff --git a/depthfm/dfm.py b/depthfm/dfm.py
index fa83a7a..63899cb 100644
--- a/depthfm/dfm.py
+++ b/depthfm/dfm.py
@@ -6,9 +6,8 @@ from torch import Tensor
 from functools import partial
 from torchdiffeq import odeint
 
-from unet import UNetModel
-from diffusers import AutoencoderKL
-
+#from unet import UNetModel
+from .unet.openaimodel import UNetModel
 
 def exists(val):
     return val is not None
diff --git a/depthfm/unet/__init__.py b/depthfm/unet/__init__.py
index b3d1ebb..fabd1bc 100644
--- a/depthfm/unet/__init__.py
+++ b/depthfm/unet/__init__.py
@@ -1,4 +1,4 @@
-import os
-import sys
-sys.path.append(os.path.dirname(os.path.abspath(__file__)))
-from openaimodel import UNetModel
\ No newline at end of file
+#import os
+#import sys
+#sys.path.append(os.path.dirname(os.path.abspath(__file__)))
+#from .openaimodel import UNetModel
\ No newline at end of file
diff --git a/depthfm/unet/attention.py b/depthfm/unet/attention.py
index b8e9f6b..1c96a47 100644
--- a/depthfm/unet/attention.py
+++ b/depthfm/unet/attention.py
@@ -6,7 +6,7 @@ from inspect import isfunction
 import torch.nn.functional as F
 from typing import Optional, Any
 
-from util import checkpoint
+from .util import checkpoint
 
 class Conv2d(torch.nn.Conv2d):
     def reset_parameters(self):
diff --git a/depthfm/unet/openaimodel.py b/depthfm/unet/openaimodel.py
index 69eadf1..50a7e67 100644
--- a/depthfm/unet/openaimodel.py
+++ b/depthfm/unet/openaimodel.py
@@ -5,7 +5,7 @@ import torch.nn as nn
 from abc import abstractmethod
 import torch.nn.functional as F
 
-from util import (
+from .util import (
     checkpoint,
     conv_nd,
     linear,
@@ -14,7 +14,7 @@ from util import (
     normalization,
     timestep_embedding,
 )
-from attention import SpatialTransformer
+from .attention import SpatialTransformer
 
 
 def exists(x):
diff --git a/nodes.py b/nodes.py
index 42c3067..6793180 100644
--- a/nodes.py
+++ b/nodes.py
@@ -1,7 +1,7 @@
 import os
 import torch
 import torch.nn.functional as F
-from .depthfm import DepthFM
+from .depthfm.dfm import DepthFM
 import folder_paths
 import comfy.utils
 import model_management