admin-valentin/WhisperLiveKit
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Place all tensors on the same device in sortformer diarization

Browse source
This commit is contained in:
Quentin Fuxa 2025-09-08 10:20:57 +02:00
parent 334b338ab0
commit 4209d7f7c0
1 changed files with 20 additions and 13 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

33
whisperlivekit/diarization/sortformer_backend.py
View file

@ -60,11 +60,15 @@ class SortformerDiarization:
self.diar_model = SortformerEncLabelModel.from_pretrained(model_name) self.diar_model = SortformerEncLabelModel.from_pretrained(model_name)
self.diar_model.eval() self.diar_model.eval()
if torch.cuda.is_available(): device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
self.diar_model.to(torch.device("cuda")) self.diar_model.to(device)
logger.info("Using CUDA for Sortformer model")
else: ## to test
logger.info("Using CPU for Sortformer model") # for name, param in self.diar_model.named_parameters():
# if param.device != device:
# raise RuntimeError(f"Parameter {name} is on {param.device} but should be on {device}")
logger.info(f"Using {device.type.upper()} for Sortformer model")
self.diar_model.sortformer_modules.chunk_len = 10 self.diar_model.sortformer_modules.chunk_len = 10
self.diar_model.sortformer_modules.subsampling_factor = 10 self.diar_model.sortformer_modules.subsampling_factor = 10
@ -187,22 +191,25 @@ class SortformerDiarizationOnline:
audio = self.buffer_audio[:threshold] audio = self.buffer_audio[:threshold]
self.buffer_audio = self.buffer_audio[threshold:] self.buffer_audio = self.buffer_audio[threshold:]
audio_signal_chunk = torch.tensor(audio).unsqueeze(0).to(self.diar_model.device) device = self.diar_model.device
audio_signal_length_chunk = torch.tensor([audio_signal_chunk.shape[1]]).to(self.diar_model.device) audio_signal_chunk = torch.tensor(audio, device=device).unsqueeze(0)
audio_signal_length_chunk = torch.tensor([audio_signal_chunk.shape[1]], device=device)
processed_signal_chunk, processed_signal_length_chunk = self.audio2mel.get_features( processed_signal_chunk, processed_signal_length_chunk = self.audio2mel.get_features(
audio_signal_chunk, audio_signal_length_chunk audio_signal_chunk, audio_signal_length_chunk
) )
processed_signal_chunk = processed_signal_chunk.to(device)
processed_signal_length_chunk = processed_signal_length_chunk.to(device)
if self._previous_chunk_features is not None: if self._previous_chunk_features is not None:
to_add = self._previous_chunk_features[:, :, -99:] to_add = self._previous_chunk_features[:, :, -99:].to(device)
total_features = torch.concat([to_add, processed_signal_chunk], dim=2) total_features = torch.concat([to_add, processed_signal_chunk], dim=2).to(device)
else: else:
total_features = processed_signal_chunk total_features = processed_signal_chunk.to(device)
self._previous_chunk_features = processed_signal_chunk self._previous_chunk_features = processed_signal_chunk.to(device)
chunk_feat_seq_t = torch.transpose(total_features, 1, 2) chunk_feat_seq_t = torch.transpose(total_features, 1, 2).to(device)
with torch.inference_mode(): with torch.inference_mode():
left_offset = 8 if self._chunk_index > 0 else 0 left_offset = 8 if self._chunk_index > 0 else 0
@ -210,7 +217,7 @@ class SortformerDiarizationOnline:
self.streaming_state, self.total_preds = self.diar_model.forward_streaming_step( self.streaming_state, self.total_preds = self.diar_model.forward_streaming_step(
processed_signal=chunk_feat_seq_t, processed_signal=chunk_feat_seq_t,
processed_signal_length=torch.tensor([chunk_feat_seq_t.shape[1]]), processed_signal_length=torch.tensor([chunk_feat_seq_t.shape[1]]).to(device),
streaming_state=self.streaming_state, streaming_state=self.streaming_state,
total_preds=self.total_preds, total_preds=self.total_preds,
left_offset=left_offset, left_offset=left_offset,