Skip to content

Fix Windows GPU diarization for YouDub #264

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
Zhuyuanbin wants to merge 1 commit into
base: main
Choose a base branch
from
Open

Fix Windows GPU diarization for YouDub #264

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
17 changes: 12 additions & 5 deletions speakerlab/bin/infer_diarization.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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.')

@gemini-code-assist gemini-code-assist Bot Jun 16, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

critical

Removing the top-level import of Inference and Model will cause a NameError in do_segmentation (line 225) because Inference is used there but is no longer in the global scope. To avoid import errors on systems where pyannote.audio is not installed while keeping the imports available globally, you can use a try...except ImportError block at the top level.

Suggested change
parser = argparse.ArgumentParser(description='Speaker diarization inference.')
try:
from pyannote.audio import Inference, Model
except ImportError:
Inference = None
Model = None
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')
Expand All @@ -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):
Expand Down Expand Up @@ -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 = {
Comment on lines +123 to 125

@gemini-code-assist gemini-code-assist Bot Jun 16, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

With Inference and Model imported at the top level using a try...except block, this local import is no longer necessary and can be removed.

    segmentation_params = {

'segmentation':'pyannote/segmentation-3.0',
'segmentation_batch_size':32,
Expand Down Expand Up @@ -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):
Expand All @@ -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)
Expand Down
17 changes: 8 additions & 9 deletions speakerlab/process/cluster.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down Expand Up @@ -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,
Expand Down
4 changes: 3 additions & 1 deletion speakerlab/utils/fileio.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,7 @@
import torch
import torchaudio
import numpy as np
import soundfile as sf


def load_yaml(yaml_path):
Expand Down Expand Up @@ -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)
Expand Down