Revert "allow cpu spk diarazation"

This reverts commit 4cb4a31ed709eb1408bfa0bf05b28f493f0a24c6.

app.py

@@ -51,13 +51,14 @@ WHISPER_LANGUAGES = [
 
 SENSEVOICE_LANGUAGES = ["auto", "zh", "yue", "en", "ja", "ko", "nospeech"]
 
-# —————— Caches & Converter ——————
+# —————— Caches ——————
 whisper_pipes = {}
 sense_models = {}
 dar_pipe = None
+
 converter = opencc.OpenCC('s2t')
 
-# —————— Helper Functions ——————
+# —————— Helpers ——————
 def get_whisper_pipe(model_id: str, device: int):
     key = (model_id, device)
     if key not in whisper_pipes:
@@ -102,14 +103,13 @@ def get_diarization_pipe():
             )
     return dar_pipe
 
+
 # —————— Whisper Transcription ——————
-def _transcribe_whisper_cpu(model_id, language, audio_path, enable_diar, diar_device):
+def _transcribe_whisper_cpu(model_id, language, audio_path, enable_diar):
     pipe = get_whisper_pipe(model_id, -1)
+    # Diarization-only branch
     if enable_diar:
         diarizer = get_diarization_pipe()
-        # Move diarization pipeline to correct device
-        dev = torch.device('cuda') if diar_device == 'GPU' and torch.cuda.is_available() else torch.device('cpu')
-        diarizer.to(dev)
         diary = diarizer(audio_path)
         snippets = []
         for turn, _, speaker in diary.itertracks(yield_label=True):
@@ -123,17 +123,18 @@ def _transcribe_whisper_cpu(model_id, language, audio_path, enable_diar, diar_de
             text = converter.convert(out.get("text", "").strip())
             snippets.append(f"[{speaker}] {text}")
         return "", "\n".join(snippets)
+    # Raw-only branch
     result = pipe(audio_path) if language == "auto" else pipe(audio_path, generate_kwargs={"language": language})
     transcript = converter.convert(result.get("text", "").strip())
     return transcript, ""
 
 
-def _transcribe_whisper_gpu(model_id, language, audio_path, enable_diar, diar_device):
+@spaces.GPU(duration=100)
+def _transcribe_whisper_gpu(model_id, language, audio_path, enable_diar):
     pipe = get_whisper_pipe(model_id, 0)
+    # Diarization-only branch
     if enable_diar:
         diarizer = get_diarization_pipe()
-        dev = torch.device('cuda') if diar_device == 'GPU' and torch.cuda.is_available() else torch.device('cpu')
-        diarizer.to(dev)
         diary = diarizer(audio_path)
         snippets = []
         for turn, _, speaker in diary.itertracks(yield_label=True):
@@ -147,23 +148,28 @@ def _transcribe_whisper_gpu(model_id, language, audio_path, enable_diar, diar_de
             text = converter.convert(out.get("text", "").strip())
             snippets.append(f"[{speaker}] {text}")
         return "", "\n".join(snippets)
+    # Raw-only branch
     result = pipe(audio_path) if language == "auto" else pipe(audio_path, generate_kwargs={"language": language})
     transcript = converter.convert(result.get("text", "").strip())
     return transcript, ""
 
 
-def transcribe_whisper(model_id, language, audio_path, device_sel, diar_device, enable_diar):
+def transcribe_whisper(model_id, language, audio_path, device_sel, enable_diar):
     if device_sel == "GPU" and torch.cuda.is_available():
-        return _transcribe_whisper_gpu(model_id, language, audio_path, enable_diar, diar_device)
-    return _transcribe_whisper_cpu(model_id, language, audio_path, enable_diar, diar_device)
+        return _transcribe_whisper_gpu(model_id, language, audio_path, enable_diar)
+    return _transcribe_whisper_cpu(model_id, language, audio_path, enable_diar)
+
 
 # —————— SenseVoice Transcription ——————
-def _transcribe_sense_cpu(model_id: str, language: str, audio_path: str, enable_punct: bool, enable_diar: bool, diar_device: str):
+def _transcribe_sense_cpu(model_id: str,
+                          language: str,
+                          audio_path: str,
+                          enable_punct: bool,
+                          enable_diar: bool):
     model = get_sense_model(model_id, "cpu")
+    # Diarization-only branch
     if enable_diar:
         diarizer = get_diarization_pipe()
-        dev = torch.device('cuda') if diar_device == 'GPU' and torch.cuda.is_available() else torch.device('cpu')
-        diarizer.to(dev)
         diary = diarizer(audio_path)
         snippets = []
         for turn, _, speaker in diary.itertracks(yield_label=True):
@@ -172,8 +178,15 @@ def _transcribe_sense_cpu(model_id: str, language: str, audio_path: str, enable_
             segment = AudioSegment.from_file(audio_path)[start_ms:end_ms]
             with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as tmp:
                 segment.export(tmp.name, format="wav")
-                segs = model.generate(input=tmp.name, cache={}, language=language, use_itn=True,
-                                      batch_size_s=300, merge_vad=False, merge_length_s=0)
+                segs = model.generate(
+                    input=tmp.name,
+                    cache={},
+                    language=language,
+                    use_itn=True,
+                    batch_size_s=300,
+                    merge_vad=False,
+                    merge_length_s=0,
+                )
             os.unlink(tmp.name)
             txt = rich_transcription_postprocess(segs[0]['text'])
             if not enable_punct:
@@ -181,8 +194,16 @@ def _transcribe_sense_cpu(model_id: str, language: str, audio_path: str, enable_
             txt = converter.convert(txt)
             snippets.append(f"[{speaker}] {txt}")
         return "", "\n".join(snippets)
-    segs = model.generate(input=audio_path, cache={}, language=language, use_itn=True,
-                          batch_size_s=300, merge_vad=True, merge_length_s=15)
+    # Raw-only branch
+    segs = model.generate(
+        input=audio_path,
+        cache={},
+        language=language,
+        use_itn=True,
+        batch_size_s=300,
+        merge_vad=True,
+        merge_length_s=15,
+    )
     text = rich_transcription_postprocess(segs[0]['text'])
     if not enable_punct:
         text = re.sub(r"[^\w\s]", "", text)
@@ -190,12 +211,16 @@ def _transcribe_sense_cpu(model_id: str, language: str, audio_path: str, enable_
     return text, ""
 
 
-def _transcribe_sense_gpu(model_id: str, language: str, audio_path: str, enable_punct: bool, enable_diar: bool, diar_device: str):
+@spaces.GPU(duration=100)
+def _transcribe_sense_gpu(model_id: str,
+                          language: str,
+                          audio_path: str,
+                          enable_punct: bool,
+                          enable_diar: bool):
     model = get_sense_model(model_id, "cuda:0")
+    # Diarization-only branch
     if enable_diar:
         diarizer = get_diarization_pipe()
-        dev = torch.device('cuda') if diar_device == 'GPU' and torch.cuda.is_available() else torch.device('cpu')
-        diarizer.to(dev)
         diary = diarizer(audio_path)
         snippets = []
         for turn, _, speaker in diary.itertracks(yield_label=True):
@@ -204,8 +229,15 @@ def _transcribe_sense_gpu(model_id: str, language: str, audio_path: str, enable_
             segment = AudioSegment.from_file(audio_path)[start_ms:end_ms]
             with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as tmp:
                 segment.export(tmp.name, format="wav")
-                segs = model.generate(input=tmp.name, cache={}, language=language, use_itn=True,
-                                      batch_size_s=300, merge_vad=False, merge_length_s=0)
+                segs = model.generate(
+                    input=tmp.name,
+                    cache={},
+                    language=language,
+                    use_itn=True,
+                    batch_size_s=300,
+                    merge_vad=False,
+                    merge_length_s=0,
+                )
             os.unlink(tmp.name)
             txt = rich_transcription_postprocess(segs[0]['text'])
             if not enable_punct:
@@ -213,8 +245,16 @@ def _transcribe_sense_gpu(model_id: str, language: str, audio_path: str, enable_
             txt = converter.convert(txt)
             snippets.append(f"[{speaker}] {txt}")
         return "", "\n".join(snippets)
-    segs = model.generate(input=audio_path, cache={}, language=language, use_itn=True,
-                          batch_size_s=300, merge_vad=True, merge_length_s=15)
+    # Raw-only branch
+    segs = model.generate(
+        input=audio_path,
+        cache={},
+        language=language,
+        use_itn=True,
+        batch_size_s=300,
+        merge_vad=True,
+        merge_length_s=15,
+    )
     text = rich_transcription_postprocess(segs[0]['text'])
     if not enable_punct:
         text = re.sub(r"[^\w\s]", "", text)
@@ -222,51 +262,67 @@ def _transcribe_sense_gpu(model_id: str, language: str, audio_path: str, enable_
     return text, ""
 
 
-def transcribe_sense(model_id: str, language: str, audio_path: str, enable_punct: bool, enable_diar: bool, device_sel: str, diar_device: str):
+def transcribe_sense(model_id: str,
+                     language: str,
+                     audio_path: str,
+                     enable_punct: bool,
+                     enable_diar: bool,
+                     device_sel: str):
     if device_sel == "GPU" and torch.cuda.is_available():
-        return _transcribe_sense_gpu(model_id, language, audio_path, enable_punct, enable_diar, diar_device)
-    return _transcribe_sense_cpu(model_id, language, audio_path, enable_punct, enable_diar, diar_device)
+        return _transcribe_sense_gpu(model_id, language, audio_path, enable_punct, enable_diar)
+    return _transcribe_sense_cpu(model_id, language, audio_path, enable_punct, enable_diar)
+
 
 # —————— Gradio UI ——————
 Demo = gr.Blocks()
 with Demo:
-    gr.Markdown("## Whisper vs. SenseVoice (Language, ASR & Diarization Devices)")
+    gr.Markdown("## Whisper vs. SenseVoice (Language, Device & Diarization with Simplified→Traditional Chinese)")
+
     audio_input = gr.Audio(sources=["upload", "microphone"], type="filepath", label="Audio Input")
+
+    # Examples
     examples = gr.Examples(
-        examples=[["interview.mp3"],["news.mp3"]],
-        inputs=[audio_input], label="Example Audio Files"
+        examples=[
+            ["interview.mp3"],
+            ["news.mp3"]
+        ],
+        inputs=[audio_input],
+        label="Example Audio Files"
     )
+
     with gr.Row():
         with gr.Column():
             gr.Markdown("### Whisper ASR")
             whisper_dd = gr.Dropdown(choices=WHISPER_MODELS, value=WHISPER_MODELS[0], label="Whisper Model")
             whisper_lang = gr.Dropdown(choices=WHISPER_LANGUAGES, value="auto", label="Whisper Language")
-            asr_device = gr.Radio(choices=["GPU","CPU"], value="GPU", label="ASR Device")
-            diar_device = gr.Radio(choices=["GPU","CPU"], value="CPU", label="Diarization Device")
+            device_radio = gr.Radio(choices=["GPU", "CPU"], value="GPU", label="Device")
             diar_check = gr.Checkbox(label="Enable Diarization", value=True)
             out_w = gr.Textbox(label="Transcript", visible=False)
             out_w_d = gr.Textbox(label="Diarized Transcript", visible=True)
+            # Toggle visibility based on checkbox
             diar_check.change(lambda e: gr.update(visible=not e), inputs=diar_check, outputs=out_w)
             diar_check.change(lambda e: gr.update(visible=e), inputs=diar_check, outputs=out_w_d)
             btn_w = gr.Button("Transcribe with Whisper")
             btn_w.click(fn=transcribe_whisper,
-                        inputs=[whisper_dd,whisper_lang,audio_input,asr_device,diar_device,diar_check],
-                        outputs=[out_w,out_w_d])
+                        inputs=[whisper_dd, whisper_lang, audio_input, device_radio, diar_check],
+                        outputs=[out_w, out_w_d])
+
         with gr.Column():
             gr.Markdown("### FunASR SenseVoice ASR")
             sense_dd = gr.Dropdown(choices=SENSEVOICE_MODELS, value=SENSEVOICE_MODELS[0], label="SenseVoice Model")
             sense_lang = gr.Dropdown(choices=SENSEVOICE_LANGUAGES, value="auto", label="SenseVoice Language")
-            asr_device_s = gr.Radio(choices=["GPU","CPU"], value="GPU", label="ASR Device")
-            diar_device_s = gr.Radio(choices=["GPU","CPU"], value="CPU", label="Diarization Device")
+            device_radio_sense = gr.Radio(choices=["GPU", "CPU"], value="GPU", label="Device")
             punct_chk = gr.Checkbox(label="Enable Punctuation", value=True)
             diar_s_chk = gr.Checkbox(label="Enable Diarization", value=True)
             out_s = gr.Textbox(label="Transcript", visible=False)
             out_s_d = gr.Textbox(label="Diarized Transcript", visible=True)
+            # Toggle visibility
             diar_s_chk.change(lambda e: gr.update(visible=not e), inputs=diar_s_chk, outputs=out_s)
             diar_s_chk.change(lambda e: gr.update(visible=e), inputs=diar_s_chk, outputs=out_s_d)
             btn_s = gr.Button("Transcribe with SenseVoice")
             btn_s.click(fn=transcribe_sense,
-                        inputs=[sense_dd,sense_lang,audio_input,punct_chk,diar_s_chk,asr_device_s,diar_device_s],
-                        outputs=[out_s,out_s_d])
+                        inputs=[sense_dd, sense_lang, audio_input, punct_chk, diar_s_chk, device_radio_sense],
+                        outputs=[out_s, out_s_d])
+
 if __name__ == "__main__":
     Demo.launch()