wire fast captioning (CLAP+Whisper+VAD) into training, add LM caption checkbox

app.py

@@ -538,7 +538,7 @@ def gradio_main():
         return "\n".join(lines)
 
     # -- Training generator (direct integration, no subprocess) --
-    def train_lora_ui(audio_files, lora_name, epochs, lr, rank):
+    def train_lora_ui(audio_files, lora_name, epochs, lr, rank, use_lm_caption):
         """Generator that yields (train_log, train_btn_update, cancel_btn_update)."""
         import gc as _gc
 
@@ -642,14 +642,9 @@ def gradio_main():
              f"Epochs: {epochs} | LR: {lr} | Rank: {rank}")
         yield _log_text(), gr.Button(visible=False), gr.Button(visible=True), gr.File()
 
-        # Caption audio files: GGUF LM if ace-server running, else librosa
-        use_understand = _server_ok()
-        method = "GGUF LM (BPM, key, mood, lyrics)" if use_understand else "librosa (BPM only)"
-        _log(f"[INFO] Auto-captioning via {method}...")
-        yield _log_text(), gr.Button(visible=False), gr.Button(visible=True), gr.File()
+        # Caption audio files without user-provided sidecars
+        audio_to_caption = []
         for audio_fname in sorted(os.listdir(audio_dir)):
-            if _training_cancel.is_set():
-                break
             full_path = os.path.join(audio_dir, audio_fname)
             if not os.path.isfile(full_path):
                 continue
@@ -661,28 +656,78 @@ def gradio_main():
             sidecar_txt = os.path.join(audio_dir, stem + ".txt")
             if os.path.isfile(sidecar_json) or os.path.isfile(sidecar_txt):
                 _log(f"  {audio_fname}: using caption file")
-                yield _log_text(), gr.Button(visible=False), gr.Button(visible=True), gr.File()
                 continue
-            caption_data = None
-            if use_understand:
-                _log(f"  {audio_fname}: GGUF LM captioning...")
+            audio_to_caption.append((audio_fname, full_path, sidecar_json))
+        yield _log_text(), gr.Button(visible=False), gr.Button(visible=True), gr.File()
+
+        if audio_to_caption and use_lm_caption and _server_ok():
+            # --- Mode: GGUF LM captioning (slow, best quality) ---
+            est_total = int(total_dur * 7 + len(audio_to_caption) * 600)
+            _log(f"[INFO] LM captioning {len(audio_to_caption)} files "
+                 f"(estimated ~{est_total // 60} min)...")
+            yield _log_text(), gr.Button(visible=False), gr.Button(visible=True), gr.File()
+            for audio_fname, full_path, sidecar_json in audio_to_caption:
+                if _training_cancel.is_set():
+                    break
+                file_dur = _lr.get_duration(path=full_path)
+                file_timeout = int(file_dur * 7 + 600)
+                _log(f"  {audio_fname}: LM captioning (timeout {file_timeout // 60} min)...")
                 yield _log_text(), gr.Button(visible=False), gr.Button(visible=True), gr.File()
                 caption_data = _caption_via_understand(
-                    full_path, timeout=600,
+                    full_path, timeout=file_timeout,
                     cancel_check=lambda: _training_cancel.is_set(),
                 )
-                if not caption_data:
-                    use_understand = False
-                    _log(f"  {audio_fname}: GGUF LM too slow, skipping (preprocessing will analyze)")
-            if caption_data:
-                bpm_s = caption_data.get("bpm", "?")
-                key_s = caption_data.get("keyscale", caption_data.get("key", "?"))
-                _log(f"  {audio_fname}: OK (BPM={bpm_s}, key={key_s})")
-                with open(sidecar_json, "w") as cj:
-                    json.dump(caption_data, cj)
-            else:
-                _log(f"  {audio_fname}: will be analyzed in preprocessing (BPM + key + caption)")
+                if caption_data:
+                    bpm_s = caption_data.get("bpm", "?")
+                    key_s = caption_data.get("keyscale", caption_data.get("key", "?"))
+                    _log(f"  {audio_fname}: OK (BPM={bpm_s}, key={key_s})")
+                    with open(sidecar_json, "w") as cj:
+                        json.dump(caption_data, cj)
+                else:
+                    _log(f"  {audio_fname}: LM failed, will use fast captioning")
+                yield _log_text(), gr.Button(visible=False), gr.Button(visible=True), gr.File()
+
+        elif audio_to_caption:
+            # --- Mode: Fast captioning (CLAP + Whisper + librosa) ---
+            _log(f"[INFO] Fast captioning {len(audio_to_caption)} files "
+                 f"(CLAP tags + lyrics + BPM)...")
             yield _log_text(), gr.Button(visible=False), gr.Button(visible=True), gr.File()
+            try:
+                from caption_fast import caption_audio, unload_caption_models
+                for audio_fname, full_path, sidecar_json in audio_to_caption:
+                    if _training_cancel.is_set():
+                        break
+                    _log(f"  {audio_fname}: analyzing...")
+                    yield _log_text(), gr.Button(visible=False), gr.Button(visible=True), gr.File()
+                    try:
+                        result = caption_audio(full_path)
+                        _log(f"  {audio_fname}: {result.get('caption', '')[:60]}")
+                        if result.get("lyrics") and result["lyrics"] != "[Instrumental]":
+                            _log(f"  {audio_fname}: lyrics extracted ({len(result['lyrics'])} chars)")
+                        with open(sidecar_json, "w") as cj:
+                            json.dump(result, cj)
+                    except Exception as cap_exc:
+                        _log(f"  {audio_fname}: fast caption failed: {cap_exc}")
+                    yield _log_text(), gr.Button(visible=False), gr.Button(visible=True), gr.File()
+                unload_caption_models()
+                _gc.collect()
+            except ImportError:
+                _log("[WARN] Fast captioning not available, using librosa fallback")
+                for audio_fname, full_path, sidecar_json in audio_to_caption:
+                    try:
+                        y_cap, sr_cap = _lr.load(full_path, sr=None, mono=True)
+                        tempo_arr, _ = _lr.beat.beat_track(y=y_cap, sr=sr_cap)
+                        bpm_val = int(round(float(
+                            tempo_arr.item() if hasattr(tempo_arr, 'item') else tempo_arr)))
+                        fallback = {"caption": audio_fname.rsplit(".", 1)[0],
+                                    "bpm": str(bpm_val), "key": "", "signature": "4/4",
+                                    "lyrics": "[Instrumental]"}
+                        with open(sidecar_json, "w") as cj:
+                            json.dump(fallback, cj)
+                        _log(f"  {audio_fname}: librosa BPM={bpm_val}")
+                    except Exception as exc:
+                        _log(f"  {audio_fname}: failed: {exc}")
+                yield _log_text(), gr.Button(visible=False), gr.Button(visible=True), gr.File()
 
         if _training_cancel.is_set():
             _training_cancel.clear()
@@ -952,6 +997,10 @@ def gradio_main():
                             label="Rank (r)", minimum=1, maximum=128,
                             value=16, step=1,
                         )
+                        use_lm_caption = gr.Checkbox(
+                            label="Use LM captioning (best quality, ~30 min/file)",
+                            value=False,
+                        )
 
                 # Button swap on click (separate handler, like rvc-beatrice)
                 # This fires immediately so user sees Cancel even if training
@@ -964,7 +1013,7 @@ def gradio_main():
                 # Training generator -- yields (log, train_btn, cancel_btn, output_file)
                 train_event = train_btn.click(
                     train_lora_ui,
-                    inputs=[train_audio, lora_name, train_epochs, train_lr, train_rank],
+                    inputs=[train_audio, lora_name, train_epochs, train_lr, train_rank, use_lm_caption],
                     outputs=[train_log, train_btn, cancel_btn, train_output_file],
                     api_name="train_lora",
                     concurrency_limit=1,