Merge pull request #8 from MomsVoiceAI/fix/answer-audio

fix: Answer audio - use Mimi streaming decoder + flash-attn

app.py

@@ -420,10 +420,21 @@ print("[MomsVoice] Creating stock voice profile (vivian)...")
 _ensure_stock_voice_profile()
 print("[MomsVoice] Stock voice profile ready.")
 
-print("[MomsVoice] Preloading LFM2.5-Audio-1.5B Q&A model...")
-from inference_lfm import get_lfm_model
-get_lfm_model()
-print("[MomsVoice] LFM Q&A model ready. All models preloaded.")
+# Load LFM Q&A model asynchronously — users typically start with Play tab,
+# so LFM can load in the background while they browse stories.
+import threading
+
+def _load_lfm_background():
+    from inference_lfm import get_lfm_model, warmup_lfm
+    print("[MomsVoice] (background) Loading LFM2.5-Audio-1.5B Q&A model...")
+    get_lfm_model()
+    print("[MomsVoice] (background) LFM loaded. Running warmup generation...")
+    warmup_lfm()
+    print("[MomsVoice] (background) LFM Q&A model ready (compiled + warmed up).")
+
+_lfm_thread = threading.Thread(target=_load_lfm_background, daemon=True)
+_lfm_thread.start()
+print("[MomsVoice] LFM loading in background thread. App starting now.")
 
 
 # Gradio Application Core setup

inference_lfm.py

@@ -70,6 +70,16 @@ def _module_device(module, fallback: torch.device) -> torch.device:
         return getattr(module, "device", fallback)
 
 
+def _reset_mimi_streaming_state(mimi) -> None:
+    """Best-effort reset of Mimi's internal streaming buffers."""
+    for attr in ("_streaming_state", "_states"):
+        if hasattr(mimi, attr):
+            try:
+                setattr(mimi, attr, None)
+            except Exception:
+                pass
+
+
 def get_lfm_model():
     """Load LFM2.5-Audio-1.5B model. Cached after first call."""
     global _processor, _model
@@ -98,21 +108,59 @@ def get_lfm_model():
                 if moved_model is not None:
                     _model = moved_model
                 _model = _model.eval()
+                # torch.compile for inference speedup on Linux/CUDA (skip on Windows)
+                if device.type == "cuda" and os.name != "nt":
+                    try:
+                        _model = torch.compile(_model, mode="reduce-overhead", dynamic=True)
+                        logger.info("torch.compile applied to LFM model (reduce-overhead, dynamic).")
+                    except Exception as exc:
+                        logger.warning("torch.compile failed, using eager mode: %s", exc)
                 # Force lazy Mimi construction after the processor is on the target device,
                 # and fail early if the streaming decoder cannot run there.
                 mimi = _processor.mimi.eval()
-                if device.type == "cuda":
-                    with torch.no_grad(), mimi.streaming(1):
-                        mimi.decode(torch.randint(0, 2048, (1, 8, 1), device=device))
+                _reset_mimi_streaming_state(mimi)
                 logger.info("LFM2.5-Audio loaded on %s.", _module_device(_model, device))
     return _processor, _model
 
 
+def warmup_lfm():
+    """Run a dummy generation to trigger torch.compile and warm CUDA graphs.
+    Call once at startup after get_lfm_model() to eliminate first-query latency.
+    """
+    try:
+        processor, model = get_lfm_model()
+        from liquid_audio import ChatState
+        device = _module_device(model, _select_device())
+        dtype = next(model.parameters()).dtype
+        mimi = processor.mimi.eval()
+
+        _reset_mimi_streaming_state(mimi)
+        with GPU_INFERENCE_LOCK, torch.no_grad():
+            chat = ChatState(processor, dtype=dtype)
+            chat.new_turn("system")
+            chat.add_text("You are a helpful assistant.")
+            chat.end_turn()
+            chat.new_turn("user")
+            chat.add_text("Hi")
+            chat.end_turn()
+            chat.new_turn("assistant")
+
+            with mimi.streaming(1):
+                for t in model.generate_interleaved(
+                    **chat, max_new_tokens=10, audio_temperature=1.0, audio_top_k=4,
+                ):
+                    pass  # Just trigger compilation, discard output
+        _reset_mimi_streaming_state(mimi)
+        logger.info("LFM warmup complete — torch.compile graphs cached.")
+    except Exception as exc:
+        logger.warning("LFM warmup failed (non-fatal): %s", exc)
+
+
 def answer_question_audio(
     question_audio_path: str | None = None,
     question_text: str | None = None,
     story_context: str = "",
-    max_new_tokens: int = 200,
+    max_new_tokens: int = 100,
 ) -> tuple[str, np.ndarray | None, int]:
     """
     Answer a question about the story using LFM2.5-Audio end-to-end.