add spaces annotation

modules/diarize/diarizer.py

@@ -3,12 +3,14 @@ import torch
 from typing import List
 import time
 import logging
+import spaces
 
 from modules.diarize.diarize_pipeline import DiarizationPipeline, assign_word_speakers
 from modules.diarize.audio_loader import load_audio
 
 
 class Diarizer:
+    @spaces.GPU
     def __init__(self,
                  model_dir: str = os.path.join("models", "Diarization")
                  ):
@@ -19,6 +21,7 @@ class Diarizer:
         os.makedirs(self.model_dir, exist_ok=True)
         self.pipe = None
 
+    @spaces.GPU
     def run(self,
             audio: str,
             transcribed_result: List[dict],
@@ -73,6 +76,7 @@ class Diarizer:
         elapsed_time = time.time() - start_time
         return diarized_result["segments"], elapsed_time
 
+    @spaces.GPU
     def update_pipe(self,
                     use_auth_token: str,
                     device: str
@@ -110,6 +114,7 @@ class Diarizer:
         logger.disabled = False
 
     @staticmethod
+    @spaces.GPU
     def get_device():
         if torch.cuda.is_available():
             return "cuda"
@@ -119,6 +124,7 @@ class Diarizer:
             return "cpu"
 
     @staticmethod
+    @spaces.GPU
     def get_available_device():
         devices = ["cpu"]
         if torch.cuda.is_available():

modules/translation/nllb_inference.py

@@ -1,11 +1,13 @@
 from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, pipeline
 import gradio as gr
 import os
+import spaces
 
 from modules.translation.translation_base import TranslationBase
 
 
 class NLLBInference(TranslationBase):
+    @spaces.GPU
     def __init__(self,
                  model_dir: str,
                  output_dir: str
@@ -20,12 +22,14 @@ class NLLBInference(TranslationBase):
         self.available_target_langs = list(NLLB_AVAILABLE_LANGS.keys())
         self.pipeline = None
 
+    @spaces.GPU
     def translate(self,
                   text: str
                   ):
         result = self.pipeline(text)
         return result[0]['translation_text']
 
+    @spaces.GPU
     def update_model(self,
                      model_size: str,
                      src_lang: str,

modules/vad/silero_vad.py

@@ -4,12 +4,14 @@ from typing import BinaryIO, Union, List, Optional
 import warnings
 import faster_whisper
 import gradio as gr
+import spaces
 
 
 class SileroVAD:
     def __init__(self):
         self.sampling_rate = 16000
 
+    @spaces.GPU
     def run(self,
             audio: Union[str, BinaryIO, np.ndarray],
             vad_parameters: VadOptions,
@@ -55,6 +57,7 @@ class SileroVAD:
         return audio
 
     @staticmethod
+    @spaces.GPU
     def get_speech_timestamps(
         audio: np.ndarray,
         vad_options: Optional[VadOptions] = None,

modules/whisper/whisper_base.py

@@ -9,6 +9,7 @@ from datetime import datetime
 from argparse import Namespace
 from faster_whisper.vad import VadOptions
 from dataclasses import astuple
+import spaces
 
 from modules.utils.subtitle_manager import get_srt, get_vtt, get_txt, write_file, safe_filename
 from modules.utils.youtube_manager import get_ytdata, get_ytaudio
@@ -18,6 +19,7 @@ from modules.vad.silero_vad import SileroVAD
 
 
 class WhisperBase(ABC):
+    @spaces.GPU
     def __init__(self,
                  model_dir: str,
                  output_dir: str,
@@ -41,6 +43,7 @@ class WhisperBase(ABC):
         self.vad = SileroVAD()
 
     @abstractmethod
+    @spaces.GPU
     def transcribe(self,
                    audio: Union[str, BinaryIO, np.ndarray],
                    progress: gr.Progress,
@@ -49,6 +52,7 @@ class WhisperBase(ABC):
         pass
 
     @abstractmethod
+    @spaces.GPU
     def update_model(self,
                      model_size: str,
                      compute_type: str,
@@ -56,6 +60,7 @@ class WhisperBase(ABC):
                      ):
         pass
 
+    @spaces.GPU
     def run(self,
             audio: Union[str, BinaryIO, np.ndarray],
             progress: gr.Progress,
@@ -121,6 +126,7 @@ class WhisperBase(ABC):
             elapsed_time += elapsed_time_diarization
         return result, elapsed_time
 
+    @spaces.GPU
     def transcribe_file(self,
                         files: list,
                         file_format: str,
@@ -191,6 +197,7 @@ class WhisperBase(ABC):
             if not files:
                 self.remove_input_files([file.name for file in files])
 
+    @spaces.GPU
     def transcribe_mic(self,
                        mic_audio: str,
                        file_format: str,
@@ -402,6 +409,7 @@ class WhisperBase(ABC):
             return "cpu"
 
     @staticmethod
+    @spaces.GPU
     def release_cuda_memory():
         if torch.cuda.is_available():
             torch.cuda.empty_cache()

requirements.txt

@@ -5,4 +5,5 @@ faster-whisper==1.0.2
 transformers
 gradio
 pytube
-pyannote.audio==3.3.1
+pyannote.audio==3.3.1
+spaces