Update app.py

app.py

@@ -79,6 +79,10 @@ def get_mask_interval(transcribe_state, word_span):
         end = float(data[e][0]) if e < len(data) else float(data[-1][1])
 
     return (start, end)
+
+
+from whisperx import load_align_model
+
     
 @spaces.GPU
 class WhisperxAlignModel:
@@ -91,22 +95,6 @@ class WhisperxAlignModel:
         audio = load_audio(audio_path)
         return align(segments, self.model, self.metadata, audio, device, return_char_alignments=False)["segments"]
 
-@spaces.GPU
-class WhisperModel:
-    def __init__(self, model_name, language):
-        from whisper import load_model
-        self.model = load_model(model_name, device, language=language)
-
-        from whisper.tokenizer import get_tokenizer
-        tokenizer = get_tokenizer(multilingual=False, language=language)
-        self.supress_tokens = [-1] + [
-            i
-            for i in range(tokenizer.eot)
-            if all(c in "0123456789" for c in tokenizer.decode([i]).removeprefix(" "))
-        ]
-
-    def transcribe(self, audio_path):
-        return self.model.transcribe(audio_path, suppress_tokens=self.supress_tokens, word_timestamps=True)["segments"]
 
 @spaces.GPU
 class WhisperxModel:
@@ -121,41 +109,45 @@ class WhisperxModel:
             segment['text'] = replace_numbers_with_words(segment['text'])
         return self.align_model.align(segments, audio_path)
 
-@spaces.GPU
-def load_models():
-    ssrspeech_model_name = "English"
-    text_tokenizer = TextTokenizer(backend="espeak")
-    language = "en"
-    transcribe_model_name = "base.en"
-
-    align_model = WhisperxAlignModel(language)
-    transcribe_model = WhisperxModel(transcribe_model_name, align_model, language)
-
-    ssrspeech_fn = f"{MODELS_PATH}/{ssrspeech_model_name}.pth"
-    if not os.path.exists(ssrspeech_fn):
-        os.system(f"wget https://huggingface.co/westbrook/SSR-Speech-{ssrspeech_model_name}/resolve/main/{ssrspeech_model_name}.pth -O " + ssrspeech_fn)
-
-    ckpt = torch.load(ssrspeech_fn)
-    model = ssr.SSR_Speech(ckpt["config"])
-    model.load_state_dict(ckpt["model"])
-    config = model.args
-    phn2num = ckpt["phn2num"]
-    model.to(device)
-
-    encodec_fn = f"{MODELS_PATH}/wmencodec.th"
-    if not os.path.exists(encodec_fn):
-        os.system(f"wget https://huggingface.co/westbrook/SSR-Speech-English/resolve/main/wmencodec.th -O " + encodec_fn)
-
-    ssrspeech_model = {
-        "config": config,
-        "phn2num": phn2num,
-        "model": model,
-        "text_tokenizer": text_tokenizer,
-        "audio_tokenizer": AudioTokenizer(signature=encodec_fn)
-    }
-    return transcribe_model, align_model, ssrspeech_model
+from whisperx import load_align_model, load_model, load_audio
+from whisperx import align as align_func
+
+
+ssrspeech_model_name = "English"
+text_tokenizer = TextTokenizer(backend="espeak")
+language = "en"
+transcribe_model_name = "base.en"
+
+# align_model = WhisperxAlignModel(language)
+# transcribe_model = WhisperxModel(transcribe_model_name, align_model, language)
+
+# align_model, align_model_metadata = load_align_model(language_code=language, device=device)
+# transcribe_model = load_model(transcribe_model_name, device, asr_options={"suppress_numerals": True, "max_new_tokens": None, "clip_timestamps": None, "hallucination_silence_threshold": None}, language=language)
+
+ssrspeech_fn = f"{MODELS_PATH}/{ssrspeech_model_name}.pth"
+if not os.path.exists(ssrspeech_fn):
+    os.system(f"wget https://huggingface.co/westbrook/SSR-Speech-{ssrspeech_model_name}/resolve/main/{ssrspeech_model_name}.pth -O " + ssrspeech_fn)
+
+ckpt = torch.load(ssrspeech_fn)
+model = ssr.SSR_Speech(ckpt["config"])
+model.load_state_dict(ckpt["model"])
+config = model.args
+phn2num = ckpt["phn2num"]
+model.to(device)
+
+encodec_fn = f"{MODELS_PATH}/wmencodec.th"
+if not os.path.exists(encodec_fn):
+    os.system(f"wget https://huggingface.co/westbrook/SSR-Speech-English/resolve/main/wmencodec.th -O " + encodec_fn)
+
+ssrspeech_model = {
+    "config": config,
+    "phn2num": phn2num,
+    "model": model,
+    "text_tokenizer": text_tokenizer,
+    "audio_tokenizer": AudioTokenizer(signature=encodec_fn)
+}
+
 
-transcribe_model, align_model, ssrspeech_model = load_models()
 
 def get_transcribe_state(segments):
     transcript = " ".join([segment["text"] for segment in segments])
@@ -167,9 +159,8 @@ def get_transcribe_state(segments):
 
 @spaces.GPU
 def transcribe(audio_path):
-    if transcribe_model is None:
-        raise gr.Error("Transcription model not loaded")
-
+    align_model, _ = load_align_model(language_code=language, device=device)
+    transcribe_model = load_model(transcribe_model_name, device, asr_options={"suppress_numerals": True, "max_new_tokens": None, "clip_timestamps": None, "hallucination_silence_threshold": None}, language=language)
     segments = transcribe_model.transcribe(audio_path)
     state = get_transcribe_state(segments)
     success_message = "<span style='color:green;'>Success: Transcribe completed successfully!</span>"
@@ -182,10 +173,9 @@ def transcribe(audio_path):
 
 @spaces.GPU
 def align(segments, audio_path):
-    if align_model is None:
-        raise gr.Error("Align model not loaded")
-    
-    segments = align_model.align(segments, audio_path)
+    align_model, metadata = load_align_model(language_code=language, device=device)
+    audio = load_audio(audio_path)
+    segments = align_func(segments, align_model, metadata, audio, device, return_char_alignments=False)["segments"]
     state = get_transcribe_state(segments)
 
     return state