diff --git a/fastmlx/utils.py b/fastmlx/utils.py
index 15976ac8..c86f2b39 100644
--- a/fastmlx/utils.py
+++ b/fastmlx/utils.py
@@ -24,14 +24,15 @@
     from mlx_lm import models as lm_models
     from mlx_lm.sample_utils import make_sampler
     from mlx_lm.tokenizer_utils import TokenizerWrapper
-    from mlx_lm.utils import generate_step
-    from mlx_lm.utils import stream_generate as lm_stream_generate
+    from mlx_lm.generate import generate_step
+    from mlx_lm.generate import stream_generate as lm_stream_generate
     from mlx_vlm import load as vlm_load
     from mlx_vlm import models as vlm_models
     from mlx_vlm.utils import load_image_processor
     from mlx_vlm.utils import stream_generate as vlm_stream_generate
-except ImportError:
+except ImportError as e:
     print("Warning: mlx or mlx_lm not available. Some functionality will be limited.")
+    print("ImportError details:", e)
 
 
 TOOLS_PATH = os.path.abspath(os.path.join(os.path.dirname(__file__), "tools"))