diff --git a/speakerlab/bin/infer_diarization.py b/speakerlab/bin/infer_diarization.py
index 2288cebae..d7b1105a8 100644
--- a/speakerlab/bin/infer_diarization.py
+++ b/speakerlab/bin/infer_diarization.py
@@ -37,8 +37,6 @@
 from modelscope.pipelines import pipeline
 from modelscope.utils.constant import Tasks
 
-from pyannote.audio import Inference, Model
-
 parser = argparse.ArgumentParser(description='Speaker diarization inference.')
 parser.add_argument('--wav', type=str, required=True, help='Input wavs')
 parser.add_argument('--out_dir', type=str, required=True, help='Out results dir')
@@ -48,6 +46,7 @@
 parser.add_argument('--diable_progress_bar', action='store_true', help='Close the progress bar')
 parser.add_argument('--nprocs', default=None, type=int, help='Num of procs')
 parser.add_argument('--speaker_num', default=None, type=int, help='Oracle num of speaker')
+parser.add_argument('--model_cache_dir', default=None, type=str, help='Local ModelScope cache root')
 
 
 def get_speaker_embedding_model(device:torch.device = None, cache_dir:str = None):
@@ -121,6 +120,8 @@ def get_voice_activity_detection_model(device: torch.device=None, cache_dir:str
     return vad_pipeline
 
 def get_segmentation_model(use_auth_token, device: torch.device=None):
+    from pyannote.audio import Inference, Model
+
     segmentation_params = {
         'segmentation':'pyannote/segmentation-3.0',
         'segmentation_batch_size':32,
@@ -424,7 +425,13 @@ def main_process(rank, nprocs, args, wav_list):
     else:
         ngpus = torch.cuda.device_count()
         device = torch.device('cuda:%d'%(rank%ngpus))
-    diarization = Diarization3Dspeaker(device, args.include_overlap, args.hf_access_token, args.speaker_num)
+    diarization = Diarization3Dspeaker(
+        device,
+        args.include_overlap,
+        args.hf_access_token,
+        args.speaker_num,
+        args.model_cache_dir,
+    )
     
     wav_list = wav_list[rank::nprocs]
     if rank == 0 and (not args.diable_progress_bar):
@@ -444,8 +451,8 @@ def main():
     if args.include_overlap and args.hf_access_token is None:
         parser.error("--hf_access_token is required when --include_overlap is specified.")
     
-    get_speaker_embedding_model()
-    get_voice_activity_detection_model()
+    get_speaker_embedding_model(cache_dir=args.model_cache_dir)
+    get_voice_activity_detection_model(cache_dir=args.model_cache_dir)
     get_cluster_backend()
     if args.include_overlap:
         get_segmentation_model(args.hf_access_token)
diff --git a/speakerlab/process/cluster.py b/speakerlab/process/cluster.py
index 1bf5bcbd7..3464953f5 100644
--- a/speakerlab/process/cluster.py
+++ b/speakerlab/process/cluster.py
@@ -11,15 +11,6 @@
 from scipy.cluster.hierarchy import fcluster
 from scipy.spatial.distance import squareform
 
-try:
-    import umap, hdbscan
-except ImportError:
-    raise ImportError(
-        "Package \"umap\" or \"hdbscan\" not found. \
-        Please install them first by \"pip install umap-learn hdbscan\"."
-        )
-
-
 class SpectralCluster:
     """A spectral clustering method using unnormalized Laplacian of affinity matrix.
     This implementation is adapted from https://github.com/speechbrain/speechbrain.
@@ -129,6 +120,14 @@ def __init__(self, n_neighbors=20, n_components=60, min_samples=20, min_cluster_
         self.metric = metric
 
     def __call__(self, X, **kwargs):
+        try:
+            import umap, hdbscan
+        except ImportError as e:
+            raise ImportError(
+                'Package "umap" or "hdbscan" not found. '
+                'Please install them first by "pip install umap-learn hdbscan".'
+            ) from e
+
         umap_X = umap.UMAP(
             n_neighbors=self.n_neighbors,
             min_dist=0.0,
diff --git a/speakerlab/utils/fileio.py b/speakerlab/utils/fileio.py
index e1fca30e1..d3c570207 100644
--- a/speakerlab/utils/fileio.py
+++ b/speakerlab/utils/fileio.py
@@ -8,6 +8,7 @@
 import torch
 import torchaudio
 import numpy as np
+import soundfile as sf
 
 
 def load_yaml(yaml_path):
@@ -104,7 +105,8 @@ def write_trans7time_list(fpath, trans7time_list):
 
 def load_audio(input, ori_fs=None, obj_fs=None):            
     if isinstance(input, str):
-        wav, fs = torchaudio.load(input)
+        wav_np, fs = sf.read(input, dtype='float32', always_2d=True)
+        wav = torch.from_numpy(wav_np.T)
         wav = wav.mean(dim=0, keepdim=True)
         if obj_fs is not None and fs != obj_fs:
             wav = torchaudio.functional.resample(wav, orig_freq=fs, new_freq=obj_fs)