perf: improve LFM2.5-Audio inference quality and reliability

Applies fixes from inference_lfm_fix.py based on demo testing:

- Fix system prompt: use official "Respond with interleaved text and audio."
instead of custom storyteller prompt; move story context to user turn
so the model processes it correctly for interleaved speech-to-speech
- Fix audio decode: switch from Mimi streaming decode to post-generation
processor.decode(audio_codes) which is more stable and avoids Mimi
streaming state corruption between calls
- Remove torch.compile and _reset_mimi_streaming_state (caused silent
audio / Mimi state bugs); replace with eager Mimi warmup on CUDA
- Raise max_new_tokens 100 -> 512 across inference_lfm, qa_flow, app.py
(100 was insufficient for complete audio generation at 12.5 Hz/frame)
- Add get_model_status() for GPU memory diagnostics
- Add text_to_audio_lfm() for LFM-based TTS mode
- Add _assemble_waveform() helper shared by TTS and Q&A paths

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

app.py

@@ -1078,7 +1078,7 @@ with gr.Blocks(title="MomsVoice", css=css_code) as demo:
             paragraphs=paragraphs,
             output_dir=SAMPLE_SOUNDS_DIR,
             answer_fn=answer_question_audio,
-            max_new_tokens=100,
+            max_new_tokens=512,
         )
 
         if not result["ok"]:

inference_lfm.py

@@ -9,6 +9,7 @@ from __future__ import annotations
 import logging
 import os
 import threading
+import time
 
 if os.name == "nt":
     os.environ.setdefault("TORCHDYNAMO_DISABLE", "1")  # Windows typically lacks Triton for torch.compile
@@ -70,14 +71,55 @@ 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):
+def _assemble_waveform(wav_chunks: list) -> np.ndarray | None:
+    """Concatenate Mimi output chunks, normalize to peak 0.9, and return float32 array."""
+    if not wav_chunks:
+        return None
+    try:
+        waveform = torch.cat(wav_chunks, dim=-1).float().numpy().squeeze()
+    except Exception as exc:
+        logger.warning("Waveform assembly failed: %s", exc)
+        return None
+
+    if waveform.ndim == 0 or waveform.size == 0:
+        return None
+
+    peak = float(np.abs(waveform).max())
+    logger.info("Waveform peak amplitude: %.6f  (samples: %d)", peak, waveform.size)
+
+    if peak == 0.0:
+        logger.warning("Waveform is all zeros — model generated silence")
+        return None
+
+    return waveform * (0.9 / peak)
+
+
+def get_model_status() -> dict:
+    """Return current model load state, device, dtype, and GPU memory."""
+    status: dict = {
+        "model_loaded": _model is not None,
+        "repo": HF_REPO,
+        "device": None,
+        "dtype": None,
+        "gpu_name": None,
+        "gpu_memory_used_gb": None,
+        "gpu_memory_reserved_gb": None,
+    }
+    if _model is not None:
+        device = _module_device(_model, _select_device())
+        status["device"] = str(device)
+        try:
+            status["dtype"] = str(next(_model.parameters()).dtype).replace("torch.", "")
+        except Exception:
+            pass
+        if device.type == "cuda":
             try:
-                setattr(mimi, attr, None)
+                status["gpu_name"] = torch.cuda.get_device_name(device)
+                status["gpu_memory_used_gb"] = round(torch.cuda.memory_allocated(device) / 1e9, 2)
+                status["gpu_memory_reserved_gb"] = round(torch.cuda.memory_reserved(device) / 1e9, 2)
             except Exception:
                 pass
+    return status
 
 
 def get_lfm_model():
@@ -108,50 +150,43 @@ 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()
-                _reset_mimi_streaming_state(mimi)
+                if device.type == "cuda":
+                    with torch.no_grad(), mimi.streaming(1):
+                        mimi.decode(torch.randint(0, 2048, (1, 8, 1), device=device))
                 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.
+    """Run a dummy generation to warm the model after loading.
     Call once at startup after get_lfm_model() to eliminate first-query latency.
     """
     try:
-        processor, model = get_lfm_model()
         from liquid_audio import ChatState
+
+        processor, model = get_lfm_model()
         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.add_text("Respond with interleaved text and audio.")
             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.")
+            for _ in model.generate_interleaved(
+                **chat, max_new_tokens=10, audio_temperature=1.0, audio_top_k=4,
+            ):
+                pass  # Just trigger warmup, discard output
+
+        logger.info("LFM warmup complete.")
     except Exception as exc:
         logger.warning("LFM warmup failed (non-fatal): %s", exc)
 
@@ -160,7 +195,7 @@ def answer_question_audio(
     question_audio_path: str | None = None,
     question_text: str | None = None,
     story_context: str = "",
-    max_new_tokens: int = 100,
+    max_new_tokens: int = 512,
 ) -> tuple[str, np.ndarray | None, int]:
     """
     Answer a question about the story using LFM2.5-Audio end-to-end.
@@ -177,19 +212,18 @@ def answer_question_audio(
     with GPU_INFERENCE_LOCK, torch.no_grad():
         chat = ChatState(processor, dtype=dtype)
 
-        # System prompt with story context
+        # Official system prompt for interleaved speech-to-speech
         chat.new_turn("system")
-        system_prompt = (
-            "You are a friendly storyteller answering a child's question about a bedtime story. "
-            "Answer in 1-2 short, simple sentences using warm, age-appropriate language. "
-            "Only use information from the story context below.\n\n"
-            f"Story context:\n{story_context[:3000]}"
-        )
-        chat.add_text(system_prompt)
+        chat.add_text("Respond with interleaved text and audio.")
         chat.end_turn()
 
-        # User turn — audio or text
+        # User turn — story context as text prefix, then audio or text question
         chat.new_turn("user")
+        if story_context:
+            chat.add_text(
+                f"Story context:\n{story_context[:3000]}\n\n"
+                "Answer the question in 1-2 short sentences based only on the story."
+            )
         if question_audio_path:
             import librosa
             wav_np, sr = librosa.load(question_audio_path, sr=16000, mono=True)
@@ -201,10 +235,83 @@ def answer_question_audio(
             return "Please ask a question!", None, SAMPLE_RATE
         chat.end_turn()
 
-        # Generate answer using Mimi streaming decode (official approach)
         chat.new_turn("assistant")
-        text_out = []
+        text_out: list[torch.Tensor] = []
+        audio_out: list[torch.Tensor] = []
+        text_token_count = 0
+        t0 = time.perf_counter()
+
+        for t in model.generate_interleaved(
+            **chat,
+            max_new_tokens=max_new_tokens,
+            audio_temperature=1.0,
+            audio_top_k=4,
+        ):
+            if t.numel() == 1:
+                text_token_count += 1
+                text_out.append(t)
+            elif t.numel() == 8:
+                audio_out.append(t)
+
+        gen_time = time.perf_counter() - t0
+        audio_frame_count = max(0, len(audio_out) - 1)  # last frame is EOS
+        logger.info(
+            "Generated: %d text tokens, %d audio frames (%.1f sec), %.1f s wall",
+            text_token_count, audio_frame_count, audio_frame_count / 12.5, gen_time,
+        )
+
+    # Decode text
+    answer_text = ""
+    if text_out:
+        answer_text = "".join(processor.text.decode(t.detach().cpu()) for t in text_out).strip()
+
+    # Decode audio via processor.decode (drops EOS frame)
+    waveform = None
+    if len(audio_out) > 1:
+        audio_codes = torch.stack(audio_out[:-1], dim=1).unsqueeze(0)  # (1, 8, N)
+        raw = processor.decode(audio_codes).cpu().float()
+        waveform = raw[0].numpy()
+        peak = float(np.abs(waveform).max())
+        if peak > 0:
+            waveform = waveform * (0.9 / peak)
+        logger.info("Decoded audio: %.2f sec, peak %.4f", len(waveform) / SAMPLE_RATE, peak)
+
+    return answer_text, waveform, SAMPLE_RATE
+
+
+def text_to_audio_lfm(
+    text: str,
+    max_new_tokens: int = 1024,
+) -> tuple[np.ndarray | None, int]:
+    """Convert text to audio using LFM2.5-Audio in TTS mode.
+
+    Feeds the text as a user message with a "read aloud" system prompt,
+    then collects only the audio tokens from generate_interleaved.
+    Returns (waveform_or_None, sample_rate).
+    """
+    from liquid_audio import ChatState
+
+    if not text.strip():
+        return None, SAMPLE_RATE
+
+    processor, model = get_lfm_model()
+    device = _module_device(model, _select_device())
+    dtype = next(model.parameters()).dtype
+
+    with GPU_INFERENCE_LOCK, torch.no_grad():
+        chat = ChatState(processor, dtype=dtype)
+
+        chat.new_turn("system")
+        chat.add_text("Read the following text aloud clearly and naturally.")
+        chat.end_turn()
+
+        chat.new_turn("user")
+        chat.add_text(text)
+        chat.end_turn()
+
+        chat.new_turn("assistant")
         wav_chunks = []
+        audio_frame_count = 0
 
         mimi = processor.mimi.eval()
         mimi_device = _first_parameter_device(mimi, device)
@@ -216,46 +323,16 @@ def answer_question_audio(
                 audio_temperature=1.0,
                 audio_top_k=4,
             ):
-                if t.numel() == 1:
-                    text_out.append(t)
-                elif t.numel() == 8:
-                    # Skip EOS/invalid marker frames before they poison the stream.
-                    if (t >= 2048).any() or (t < 0).any():
+                if t.numel() == 8:
+                    if (t == 2048).any():
                         continue
-                    codes = t.reshape(1, 8, 1).to(device=mimi_device, dtype=torch.long)
+                    audio_frame_count += 1
                     try:
-                        wav_chunk = mimi.decode(codes)
+                        wav_chunk = mimi.decode(t[None, :, None].to(device=mimi_device, dtype=torch.long))[0]
                         wav_chunks.append(wav_chunk.cpu())
                     except Exception as exc:
-                        logger.warning(
-                            "Mimi decode skipped frame: %s (codes_device=%s, mimi_device=%s, shape=%s)",
-                            exc,
-                            codes.device,
-                            mimi_device,
-                            tuple(codes.shape),
-                        )
+                        logger.warning("TTS decode skipped frame: %s", exc)
 
-    # Decode text
-    answer_text = ""
-    if text_out:
-        answer_text = "".join(processor.text.decode(t.detach().cpu()) for t in text_out).strip()
-
-    # Assemble waveform from streaming chunks
-    waveform = None
-    if wav_chunks:
-        try:
-            waveform = torch.cat(wav_chunks, dim=-1).float().numpy().squeeze()
-            if waveform.ndim == 0 or len(waveform) == 0:
-                waveform = None
-            else:
-                peak = np.abs(waveform).max()
-                if peak < 0.001:
-                    waveform = None  # Silent, likely garbage
-                else:
-                    # Normalize to target peak (0.9) for consistent volume
-                    waveform = waveform * (0.9 / peak)
-        except Exception as e:
-            logger.warning("Audio decode failed: %s", e)
-            waveform = None
+        logger.info("TTS: %d audio frames (%.1f sec)", audio_frame_count, audio_frame_count / 12.5)
 
-    return answer_text, waveform, SAMPLE_RATE
+    return _assemble_waveform(wav_chunks), SAMPLE_RATE

qa_flow.py

@@ -101,7 +101,7 @@ def build_qa_response(
     output_dir: Path,
     answer_fn: Callable[..., tuple[str, object | None, int]],
     audio_writer: AudioWriter = _default_audio_writer,
-    max_new_tokens: int = 100,
+    max_new_tokens: int = 512,
 ) -> dict[str, object]:
     """Run Q&A and return display-ready data without importing Gradio."""
     q_txt = (question_text or "").strip()