Add MMS_FA forced alignment — return word timestamps per audio slot

Each generated audio now runs forced alignment (torchaudio MMS_FA) on the
cleaned text and returns [{text, start, end}] JSON — same structure as Fish
TTS word timestamps — so callers can build SRT using their own logic.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

app.py

@@ -273,6 +273,10 @@ def load_models():
     pass
 
 
+# Load MMS_FA alignment model at startup so ZeroGPU preloads it
+_get_mms_fa()
+
+
 
 def compile_model(should_compile):
     """Compile the model for faster inference."""
@@ -322,6 +326,79 @@ def do_compile():
         return gr.update(value=f"✗ Compilation failed: {str(e)}", visible=True), gr.update(interactive=True)
 
 
+def _clean_text_echo(text: str) -> str:
+    """Strip (parenthetical cues) for SRT alignment."""
+    import re
+    return re.sub(r'\([^)]*\)', '', text).strip()
+
+
+_mms_fa_bundle = None
+_mms_fa_model = None
+_mms_fa_dict = None
+_mms_fa_sample_rate = None
+
+
+def _get_mms_fa():
+    global _mms_fa_bundle, _mms_fa_model, _mms_fa_dict, _mms_fa_sample_rate
+    if _mms_fa_model is None:
+        _mms_fa_bundle = torchaudio.pipelines.MMS_FA
+        _mms_fa_model = _mms_fa_bundle.get_model(with_star=False).cuda()
+        _mms_fa_dict = _mms_fa_bundle.get_dict()
+        _mms_fa_sample_rate = _mms_fa_bundle.sample_rate
+    return _mms_fa_model, _mms_fa_dict, _mms_fa_sample_rate
+
+
+def align_words(audio_path: str, text: str) -> list:
+    """
+    Forced-align text to audio using MMS_FA.
+    Returns [{text, start, end}, ...] — same structure as Fish TTS word timestamps.
+    """
+    model, dictionary, bundle_sr = _get_mms_fa()
+
+    waveform, sr = torchaudio.load(audio_path)
+    if waveform.shape[0] > 1:
+        waveform = waveform.mean(dim=0, keepdim=True)
+    if sr != bundle_sr:
+        waveform = torchaudio.functional.resample(waveform, sr, bundle_sr)
+    waveform = waveform.cuda()
+
+    words = text.split()
+    if not words:
+        return []
+
+    try:
+        with torch.inference_mode():
+            emission, _ = model(waveform)
+
+        tokenised = []
+        kept_words = []
+        for w in words:
+            tokens = [dictionary[c] for c in w.lower() if c in dictionary]
+            if tokens:
+                tokenised.append(tokens)
+                kept_words.append(w)
+
+        if not tokenised:
+            return []
+
+        token_spans = torchaudio.functional.forced_align(emission, tokenised, blank=0)
+
+        num_frames = emission.shape[1]
+        audio_duration = waveform.shape[-1] / bundle_sr
+
+        result = []
+        for word, spans in zip(kept_words, token_spans):
+            result.append({
+                "text": word,
+                "start": round(spans[0].start / num_frames * audio_duration, 4),
+                "end": round(spans[-1].end / num_frames * audio_duration, 4),
+            })
+        return result
+    except Exception as e:
+        print(f"[align] Alignment failed: {e}")
+        return []
+
+
 def save_audio_with_format(audio_tensor: torch.Tensor, base_path: Path, filename: str, sample_rate: int, audio_format: str) -> Path:
     """Save audio in specified format, fallback to WAV if MP3 encoding fails."""
     if audio_format == "mp3":
@@ -2182,7 +2259,7 @@ def generate_unified(
     texts = texts[:MAX_TEXTS]
 
     if not texts:
-        return tuple([gr.update(visible=False)] * (1 + MAX_TEXTS * 2))
+        return tuple([gr.update(visible=False)] * (1 + MAX_TEXTS * 3))
 
     # Determine speaker audio path (URL or file)
     speaker_audio_path = None
@@ -2268,6 +2345,7 @@ def generate_unified(
     # Initialize outputs
     result_labels = [gr.update(visible=False)] * MAX_TEXTS
     result_audios = [gr.update(visible=False)] * MAX_TEXTS
+    result_srts = [gr.update(visible=False)] * MAX_TEXTS
     time_display = gr.update(visible=False)
 
     # Generate each text
@@ -2308,8 +2386,12 @@ def generate_unified(
         _cdn = upload_to_r2(audio_path)
         result_audios[i] = gr.update(value=_cdn if _cdn else str(audio_path), visible=True)
 
+        # Forced alignment → word timestamps (same structure as Fish TTS)
+        words = align_words(str(audio_path), _clean_text_echo(text))
+        result_srts[i] = gr.update(value=json.dumps(words, ensure_ascii=False), visible=True)
+
         # Yield progressive results
-        yield tuple([time_display] + result_labels + result_audios)
+        yield tuple([time_display] + result_labels + result_audios + result_srts)
 
     # Final time display
     batch_duration = time.time() - batch_start
@@ -2318,7 +2400,7 @@ def generate_unified(
         visible=True
     )
 
-    yield tuple([time_display] + result_labels + result_audios)
+    yield tuple([time_display] + result_labels + result_audios + result_srts)
 
 
 with gr.Blocks(title="Echo-TTS", css=LINK_CSS + SIMPLE_CSS, js=JS_CODE) as demo:
@@ -2426,9 +2508,11 @@ with gr.Blocks(title="Echo-TTS", css=LINK_CSS + SIMPLE_CSS, js=JS_CODE) as demo:
             # Pre-allocate output slots (up to MAX_TEXTS)
             output_labels = []
             output_audios = []
+            output_srts = []
             for i in range(MAX_TEXTS):
                 output_labels.append(gr.Markdown(visible=False))
                 output_audios.append(gr.Textbox(visible=False, interactive=False, label=f"Audio URL {i+1}"))
+                output_srts.append(gr.Textbox(visible=False, interactive=False, label=f"SRT {i+1}", lines=4))
 
     # Session state
     session_id_state = gr.State(None)
@@ -2503,7 +2587,7 @@ with gr.Blocks(title="Echo-TTS", css=LINK_CSS + SIMPLE_CSS, js=JS_CODE) as demo:
             audio_format,
             session_id_state,
         ],
-        outputs=[generation_time_display] + output_labels + output_audios
+        outputs=[generation_time_display] + output_labels + output_audios + output_srts
     )
 
     # Initialize session