Merge branch 'main' of https://huggingface.co/spaces/jhj0517/Whisper-WebUI into huggingface

modules/diarize/diarizer.py

@@ -4,6 +4,7 @@ from typing import List, Union, BinaryIO, Optional
 import numpy as np
 import time
 import logging
+import spaces
 
 from modules.utils.paths import DIARIZATION_MODELS_DIR
 from modules.diarize.diarize_pipeline import DiarizationPipeline, assign_word_speakers
@@ -21,6 +22,7 @@ class Diarizer:
         os.makedirs(self.model_dir, exist_ok=True)
         self.pipe = None
 
+    @spaces.GPU
     def run(self,
             audio: Union[str, BinaryIO, np.ndarray],
             transcribed_result: List[dict],
@@ -77,6 +79,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
@@ -115,6 +118,7 @@ class Diarizer:
         logger.disabled = False
 
     @staticmethod
+    @spaces.GPU
     def get_device():
         if torch.cuda.is_available():
             return "cuda"
@@ -124,6 +128,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,6 +1,7 @@
 from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, pipeline
 import gradio as gr
 import os
+import spaces
 
 from modules.utils.paths import TRANSLATION_OUTPUT_DIR, NLLB_MODELS_DIR
 from modules.translation.translation_base import TranslationBase
@@ -21,6 +22,7 @@ class NLLBInference(TranslationBase):
         self.available_target_langs = list(NLLB_AVAILABLE_LANGS.keys())
         self.pipeline = None
 
+    @spaces.GPU(duration=120)
     def translate(self,
                   text: str,
                   max_length: int
@@ -31,6 +33,7 @@ class NLLBInference(TranslationBase):
         )
         return result[0]['translation_text']
 
+    @spaces.GPU(duration=120)
     def update_model(self,
                      model_size: str,
                      src_lang: str,

modules/translation/translation_base.py

@@ -4,6 +4,7 @@ import gradio as gr
 from abc import ABC, abstractmethod
 from typing import List
 from datetime import datetime
+import spaces
 
 from modules.whisper.whisper_parameter import *
 from modules.utils.subtitle_manager import *
@@ -26,6 +27,7 @@ class TranslationBase(ABC):
         self.device = self.get_device()
 
     @abstractmethod
+    @spaces.GPU(duration=120)
     def translate(self,
                   text: str,
                   max_length: int
@@ -33,6 +35,7 @@ class TranslationBase(ABC):
         pass
 
     @abstractmethod
+    @spaces.GPU(duration=120)
     def update_model(self,
                      model_size: str,
                      src_lang: str,
@@ -41,6 +44,7 @@ class TranslationBase(ABC):
                      ):
         pass
 
+    @spaces.GPU(duration=120)
     def translate_file(self,
                        fileobjs: list,
                        model_size: str,
@@ -135,6 +139,7 @@ class TranslationBase(ABC):
             self.release_cuda_memory()
 
     @staticmethod
+    @spaces.GPU(duration=120)
     def get_device():
         if torch.cuda.is_available():
             return "cuda"
@@ -144,6 +149,7 @@ class TranslationBase(ABC):
             return "cpu"
 
     @staticmethod
+    @spaces.GPU(duration=120)
     def release_cuda_memory():
         if torch.cuda.is_available():
             torch.cuda.empty_cache()

modules/utils/subtitle_manager.py

@@ -1,5 +1,7 @@
 import re
 
+# Zero GPU
+import spaces
 
 def timeformat_srt(time):
     hours = time // 3600
@@ -117,7 +119,7 @@ def get_serialized_vtt(dicts):
         output += f'{dic["sentence"]}\n\n'
     return output
 
-
+@spaces.GPU(duration=120)
 def safe_filename(name):
     from app import _args
     INVALID_FILENAME_CHARS = r'[<>:"/\\|?*\x00-\x1f]'