Skip to content

Commit

Permalink
Merge pull request #244 from jhj0517/fix/device
Browse files Browse the repository at this point in the history
Fix device bug
  • Loading branch information
jhj0517 authored Aug 30, 2024
2 parents 0b0f426 + 15b3a25 commit 4e55844
Show file tree
Hide file tree
Showing 2 changed files with 8 additions and 6 deletions.
13 changes: 8 additions & 5 deletions modules/diarize/diarizer.py
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
import os
import torch
from typing import List, Union, BinaryIO
from typing import List, Union, BinaryIO, Optional
import numpy as np
import time
import logging
Expand All @@ -24,7 +24,7 @@ def run(self,
audio: Union[str, BinaryIO, np.ndarray],
transcribed_result: List[dict],
use_auth_token: str,
device: str
device: Optional[str] = None
):
"""
Diarize transcribed result as a post-processing
Expand All @@ -38,7 +38,7 @@ def run(self,
use_auth_token: str
Huggingface token with READ permission. This is only needed the first time you download the model.
You must manually go to the website https://huggingface.co/pyannote/speaker-diarization-3.1 and agree to their TOS to download the model.
device: str
device: Optional[str]
Device for diarization.
Returns
Expand All @@ -50,8 +50,10 @@ def run(self,
"""
start_time = time.time()

if (device != self.device
or self.pipe is None):
if device is None:
device = self.device

if device != self.device or self.pipe is None:
self.update_pipe(
device=device,
use_auth_token=use_auth_token
Expand Down Expand Up @@ -89,6 +91,7 @@ def update_pipe(self,
device: str
Device for diarization.
"""
self.device = device

os.makedirs(self.model_dir, exist_ok=True)

Expand Down
1 change: 0 additions & 1 deletion modules/whisper/whisper_base.py
Original file line number Diff line number Diff line change
Expand Up @@ -130,7 +130,6 @@ def run(self,
audio=audio,
use_auth_token=params.hf_token,
transcribed_result=result,
device=self.device
)
elapsed_time += elapsed_time_diarization
return result, elapsed_time
Expand Down

0 comments on commit 4e55844

Please sign in to comment.