Auto-transcribe reference audio with CPU Whisper into the editable transcript field

app.py

@@ -9,6 +9,7 @@ os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")
 
 import spaces  # MUST come before any torch/CUDA import
 import tempfile
+import threading
 import time
 from pathlib import Path
 
@@ -27,6 +28,7 @@ from breeze_infer.templates import get_template, prepare_inputs
 from models.fast_streaming import FastBreezeStreamingRuntime, FastStreamingConfig
 
 MODEL_ID = "BreezeBlue/Breeze-TTS-2"
+ASR_MODEL_ID = "openai/whisper-small"  # multilingual (en/zh), runs on CPU
 
 MAX_NEW_TOKENS = 1500
 MAX_SEQ_LEN = 2048
@@ -74,6 +76,67 @@ print("Breeze TTS 2 loaded successfully.")
 print(f"Sample rate: {_runtime.sample_rate}")
 
 
+_asr_pipe = None
+_asr_lock = threading.Lock()
+
+
+def _get_asr():
+    """Lazily build the CPU Whisper pipeline used to auto-transcribe reference audio."""
+    global _asr_pipe
+    with _asr_lock:
+        if _asr_pipe is None:
+            from transformers import pipeline
+
+            _asr_pipe = pipeline(
+                "automatic-speech-recognition",
+                model=ASR_MODEL_ID,
+                device="cpu",
+                chunk_length_s=30,
+            )
+    return _asr_pipe
+
+
+def transcribe_reference(
+    ref_audio_path: str | None,
+    progress=gr.Progress(track_tqdm=True),
+):
+    """Auto-transcribe uploaded/recorded reference audio with Whisper on CPU.
+
+    The result populates the (still editable) reference transcript field.
+
+    Args:
+        ref_audio_path: Path to the reference audio file, or None when cleared.
+    """
+    if not ref_audio_path:
+        return gr.update()
+
+    try:
+        asr = _get_asr()
+    except Exception as exc:  # pragma: no cover - model download/load failure
+        raise gr.Error(f"Could not load the transcription model: {exc}") from exc
+
+    try:
+        import librosa
+
+        audio, _ = librosa.load(ref_audio_path, sr=16000, mono=True)
+        asr_input = {"raw": audio, "sampling_rate": 16000}
+    except Exception:
+        # Fall back to letting the pipeline decode the file itself (ffmpeg).
+        asr_input = str(ref_audio_path)
+
+    try:
+        result = asr(asr_input, generate_kwargs={"task": "transcribe"})
+    except Exception as exc:
+        raise gr.Error(f"Automatic transcription failed: {exc}") from exc
+
+    text = (result.get("text") or "").strip()
+    if not text:
+        raise gr.Error(
+            "Automatic transcription produced no text — please type the transcript manually."
+        )
+    return text
+
+
 @spaces.GPU(duration=120)
 def generate_voice_design(
     text: str,
@@ -295,9 +358,16 @@ with gr.Blocks(theme=gr.themes.Citrus(), css=CSS) as demo:
                         )
                         clone_ref_text = gr.Textbox(
                             label="Reference transcript (exact)",
+                            info="Auto-transcribed with Whisper when you add reference audio — please review and edit if needed.",
                             placeholder="This is the exact transcript of the reference audio.",
                             lines=2,
                         )
+                        clone_ref_audio.change(
+                            fn=transcribe_reference,
+                            inputs=[clone_ref_audio],
+                            outputs=[clone_ref_text],
+                            api_name="transcribe_reference",
+                        )
                         with gr.Accordion("Advanced settings", open=False):
                             clone_seed = gr.Number(label="Seed", value=42, precision=0)
                         clone_btn = gr.Button("Generate", variant="primary")
@@ -338,9 +408,16 @@ with gr.Blocks(theme=gr.themes.Citrus(), css=CSS) as demo:
                         )
                         dir_ref_text = gr.Textbox(
                             label="Reference transcript (exact)",
+                            info="Auto-transcribed with Whisper when you add reference audio — please review and edit if needed.",
                             placeholder="This is the exact transcript of the reference audio.",
                             lines=2,
                         )
+                        dir_ref_audio.change(
+                            fn=transcribe_reference,
+                            inputs=[dir_ref_audio],
+                            outputs=[dir_ref_text],
+                            api_name="transcribe_reference_direction",
+                        )
                         with gr.Accordion("Advanced settings", open=False):
                             dir_cfg = gr.Slider(
                                 label="CFG scale",
@@ -378,4 +455,10 @@ with gr.Blocks(theme=gr.themes.Citrus(), css=CSS) as demo:
             "Weights/outputs are for research and non-commercial use only."
         )
 
+try:
+    _get_asr()
+    print(f"Whisper ASR ({ASR_MODEL_ID}) loaded on CPU for reference transcription.")
+except Exception as _exc:  # pragma: no cover
+    print(f"Warning: could not preload Whisper ASR ({_exc}); will retry on first use.")
+
 demo.launch(mcp_server=True)