fix: proper Mimi decode on HF Space — only disable dynamo on Windows, explicit device handling

- Only set TORCHDYNAMO_DISABLE on Windows (HF Space has Triton)
- Use torch.no_grad() matching official demo pattern
- Warmup Mimi at load time to catch device errors early
- Explicit mimi_device for code tensors
- Log decode failures with device info for debugging
- Force format=wav for mic input, restore autoplay=True

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

app.py

@@ -554,6 +554,7 @@ with gr.Blocks(title="MomsVoice", css=css_code) as demo:
                             sources=["microphone"], type="filepath",
                             label="🎤 Ask your question — answer generates when you stop recording",
                             show_label=True, streaming=False,
+                            format="wav",
                             elem_id="qa_audio_input",
                         )
                         question_text = gr.Textbox(visible=False)
@@ -1092,7 +1093,7 @@ with gr.Blocks(title="MomsVoice", css=css_code) as demo:
             </div>
         """
         if answer_audio_path_out:
-            return gr.HTML(value=answer_html, visible=True), gr.Audio(value=str(answer_audio_path_out), visible=True)
+            return gr.HTML(value=answer_html, visible=True), gr.Audio(value=str(answer_audio_path_out), visible=True, autoplay=True)
         return gr.HTML(value=answer_html, visible=True), gr.Audio(visible=False)
 
     submit_question_btn.click(

inference_lfm.py

@@ -10,7 +10,8 @@ import logging
 import os
 import threading
 
-os.environ.setdefault("TORCHDYNAMO_DISABLE", "1")  # Mimi codec uses torch.compile; skip on Windows/no-Triton
+if os.name == "nt":
+    os.environ.setdefault("TORCHDYNAMO_DISABLE", "1")  # Windows typically lacks Triton for torch.compile
 
 import torch
 import numpy as np
@@ -52,6 +53,16 @@ def _move_module(module, device: torch.device, dtype: torch.dtype):
         return module
 
 
+def _first_parameter_device(module, fallback: torch.device) -> torch.device:
+    try:
+        return next(module.parameters()).device
+    except Exception:
+        try:
+            return next(module.buffers()).device
+        except Exception:
+            return fallback
+
+
 def _module_device(module, fallback: torch.device) -> torch.device:
     try:
         return next(module.parameters()).device
@@ -70,11 +81,12 @@ def get_lfm_model():
                 device = _select_device()
                 dtype = _select_dtype()
                 logger.info("Loading %s on %s (%s)...", HF_REPO, device, dtype)
-                _processor = LFM2AudioProcessor.from_pretrained(HF_REPO).eval()
+                _processor = LFM2AudioProcessor.from_pretrained(HF_REPO, device=device).eval()
                 try:
                     _model = LFM2AudioModel.from_pretrained(
                         HF_REPO,
-                        torch_dtype=dtype if device.type == "cuda" else torch.float32,
+                        dtype=dtype if device.type == "cuda" else torch.float32,
+                        device=device,
                     ).eval()
                 except TypeError:
                     _model = LFM2AudioModel.from_pretrained(HF_REPO).eval()
@@ -86,6 +98,12 @@ def get_lfm_model():
                 if moved_model is not None:
                     _model = moved_model
                 _model = _model.eval()
+                # 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))
                 logger.info("LFM2.5-Audio loaded on %s.", _module_device(_model, device))
     return _processor, _model
 
@@ -102,63 +120,72 @@ def answer_question_audio(
     Accepts either audio input (child's voice) or text input.
     Returns (answer_text, audio_waveform_or_None, sample_rate).
     """
-    from liquid_audio import ChatState, LFMModality
+    from liquid_audio import ChatState
 
     processor, model = get_lfm_model()
-
-    chat = ChatState(processor)
-
-    # System prompt with story context
-    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.end_turn()
-
-    # User turn — audio or text
-    chat.new_turn("user")
-    if question_audio_path:
-        import librosa
-        wav_np, sr = librosa.load(question_audio_path, sr=16000, mono=True)
-        wav = torch.from_numpy(wav_np).unsqueeze(0).to(_module_device(model, _select_device()))
-        sr = 16000
-        chat.add_audio(wav, sr)
-    elif question_text:
-        chat.add_text(question_text)
-    else:
-        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 = []
-    wav_chunks = []
-
-    mimi = processor.mimi
-
-    with GPU_INFERENCE_LOCK, torch.inference_mode(), mimi.streaming(1):
-        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_out.append(t)
-            elif t.numel() == 8:
-                # Skip EOS marker (all codes == 2048)
-                if torch.all(t >= 2048):
-                    continue
-                # Decode each audio frame via Mimi streaming codec
-                try:
-                    wav_chunk = mimi.decode(t[None, :, None])
-                    wav_chunks.append(wav_chunk.cpu())
-                except Exception:
-                    pass  # Skip invalid frames
+    device = _module_device(model, _select_device())
+    dtype = next(model.parameters()).dtype
+
+    with GPU_INFERENCE_LOCK, torch.no_grad():
+        chat = ChatState(processor, dtype=dtype)
+
+        # System prompt with story context
+        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.end_turn()
+
+        # User turn — audio or text
+        chat.new_turn("user")
+        if question_audio_path:
+            import librosa
+            wav_np, sr = librosa.load(question_audio_path, sr=16000, mono=True)
+            wav = torch.from_numpy(wav_np).unsqueeze(0).to(device)
+            chat.add_audio(wav, sr)
+        elif question_text:
+            chat.add_text(question_text)
+        else:
+            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 = []
+        wav_chunks = []
+
+        mimi = processor.mimi.eval()
+        mimi_device = _first_parameter_device(mimi, device)
+
+        with mimi.streaming(1):
+            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_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():
+                        continue
+                    codes = t.reshape(1, 8, 1).to(device=mimi_device, dtype=torch.long)
+                    try:
+                        wav_chunk = mimi.decode(codes)
+                        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),
+                        )
 
     # Decode text
     answer_text = ""