"""PersonaPlex + IMTalker AJ server with split audio/video WebSockets.

Raw PersonaPlex audio is never stored in or cleared with the video queue.
Assistant Opus audio stays on /ws/conversation, while JPEG video frames are
sent through /ws/video?session_id=... to avoid websocket head-of-line blocking.
"""
from __future__ import annotations

import argparse
import asyncio
import base64
import concurrent.futures
import contextlib
import faulthandler
import json
import os
import queue
import signal

faulthandler.enable()
try:
    faulthandler.register(signal.SIGUSR1, all_threads=True)
except Exception:
    pass

import ws_av_binary_codec as _wsbin
import sys
import threading
import time
import types
import uuid
from pathlib import Path
from typing import Optional

import cv2
import numpy as np
import sphn
import torch
import torch.nn as nn
import torch.nn.functional as F
import torchaudio
import torchvision.transforms as T
from fastapi import FastAPI, WebSocket, WebSocketDisconnect
from fastapi.responses import FileResponse, HTMLResponse, JSONResponse
from fastapi.staticfiles import StaticFiles
from PIL import Image
from transformers import Wav2Vec2FeatureExtractor

ROOT = Path(__file__).resolve().parent
if str(ROOT) not in sys.path:
    sys.path.insert(0, str(ROOT))

from generator.FM import FMGenerator
from generator.train_lora import apply_lora_to_model
from generator.helium_w2v_frontend_adapter import HeliumToWav2VecFrontendAdapter
from generator.unitalk_wav2vec_adapter import UniTalkLastLayerLiveAdapter
from generator.options.base_options import BaseOptions
from generator.wav2vec2 import Wav2VecModel
if os.environ.get("IMTALKER_CACHED_ENGINE", "0").strip().lower() in {
    "1",
    "true",
    "yes",
    "on",
}:
    from liveTry_cached import MoshiOnlyEngine
else:
    from liveTry import MoshiOnlyEngine
from renderer.models import IMTRenderer

# ---------------------------------------------------------------------------
# Constants
# ---------------------------------------------------------------------------

TARGET_SR = 24_000          # Mimi sample rate (24kHz)
VIDEO_FPS = 25              # IMTalker frame rate
MIMI_FRAME_SIZE = 1_920     # samples per Mimi frame (80ms @ 24kHz)
MAIN_CODEBOOKS = 8          # codebooks used for Helium input embeddings
PREBUFFER_CHUNKS = 0        # produce this many chunks before sender starts pacing
WAV2VEC_SR = 16_000


class PlasticityProjectionHead(nn.Module):
    def __init__(self) -> None:
        super().__init__()
        self.input_ln = nn.LayerNorm(4096)
        self.net = nn.Sequential(
            nn.Linear(4096, 1024),
            nn.LayerNorm(1024),
            nn.GELU(),
            nn.Dropout(0.1),
            nn.Linear(1024, 768),
            nn.LayerNorm(768),
            nn.GELU(),
            nn.Dropout(0.1),
            nn.Linear(768, 768),
            nn.LayerNorm(768),
            nn.GELU(),
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.net(self.input_ln(x))


class PlasticityUpsampler(nn.Module):
    def __init__(self) -> None:
        super().__init__()
        self.up = nn.ConvTranspose1d(768, 768, kernel_size=4, stride=4)

    def forward(self, low: torch.Tensor, target_len: int) -> torch.Tensor:
        y = self.up(low.transpose(1, 2).contiguous())
        if y.shape[-1] != int(target_len):
            y = F.interpolate(y, size=int(target_len), mode="linear", align_corners=False)
        return y.transpose(1, 2).contiguous()


class PlasticityCausalBlock(nn.Module):
    def __init__(self) -> None:
        super().__init__()
        self.norm1 = nn.LayerNorm(768)
        self.attn = nn.MultiheadAttention(
            embed_dim=768,
            num_heads=12,
            dropout=0.15,
            batch_first=True,
        )
        self.drop1 = nn.Dropout(0.1)
        self.norm2 = nn.LayerNorm(768)
        self.ff = nn.Sequential(
            nn.Linear(768, 2048),
            nn.GELU(),
            nn.Dropout(0.1),
            nn.Linear(2048, 768),
        )
        self.drop2 = nn.Dropout(0.1)

    def forward(self, x: torch.Tensor, mask: torch.Tensor) -> torch.Tensor:
        h = self.norm1(x)
        attn, _ = self.attn(h, h, h, attn_mask=mask, need_weights=False)
        x = x + self.drop1(attn)
        h = self.norm2(x)
        x = x + self.drop2(self.ff(h))
        return x


class PlasticityCausalTransformer(nn.Module):
    def __init__(self, max_len: int = 2048) -> None:
        super().__init__()
        self.blocks = nn.ModuleList([PlasticityCausalBlock() for _ in range(8)])
        self.norm = nn.LayerNorm(768)
        mask = torch.triu(torch.ones(max_len, max_len, dtype=torch.bool), diagonal=1)
        self.register_buffer("causal_mask", mask, persistent=False)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        mask = self.causal_mask[: x.shape[1], : x.shape[1]].to(device=x.device)
        for block in self.blocks:
            x = block(x, mask)
        return self.norm(x)


class StudioNativeLiveAdapter(nn.Module):
    """Frontend fp32 adapter live wrapper.

    Training contract:
      raw 12.5Hz Helium -> Wav2Vec2 projected frontend [T50, 768]
      live contract:
      projected frontend -> frozen Wav2Vec2 encoder -> final hidden -> IMTalker audio_projection.
    """

    def __init__(self, wav2vec_model_path: str, num_layers: int, dropout: float = 0.0) -> None:
        super().__init__()
        self.model = HeliumToWav2VecFrontendAdapter(num_layers=int(num_layers), dropout=float(dropout))
        self.wav2vec = Wav2VecModel.from_pretrained(wav2vec_model_path, local_files_only=True).eval().float()
        for param in self.wav2vec.parameters():
            param.requires_grad_(False)

    def load_state_dict(self, state_dict, strict: bool = True):  # type: ignore[override]
        return self.model.load_state_dict(state_dict, strict=strict)

    @torch.no_grad()
    def forward_single(self, source: torch.Tensor, target_len: int) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
        src = source.unsqueeze(0).contiguous()
        target_len = int(target_len)
        frontend_len = max(1, target_len * 2)
        frontend50 = self.model(src.float(), target_len=frontend_len).float()
        final50 = self.wav2vec.encode_from_projected_frontend(frontend50).last_hidden_state.float()
        final25 = F.interpolate(
            final50.transpose(1, 2),
            size=target_len,
            mode="linear",
            align_corners=False,
        ).transpose(1, 2)[0].float().contiguous()
        return frontend50[0].float().contiguous(), final50[0].float().contiguous(), final25


class MoshiOnlyEngineWithHidden(MoshiOnlyEngine):
    """Moshi reply engine that also returns the main LM hidden for each generated step.

    Layer[-2] is exposed as a native LMGen output so Moshi can keep CUDA graph
    replay enabled. We do not use Python forward hooks in this path.
    """

    def __init__(self, *args, capture_layer: int = -2, **kwargs) -> None:
        self.tf_capture_layer = int(capture_layer)
        super().__init__(*args, **kwargs)
        self._install_graph_hidden_capture()

    def _install_graph_hidden_capture(self) -> None:
        lm_model = self.lm
        lm_gen = self.lm_gen
        if hasattr(lm_gen, "prepare_step_input") and hasattr(lm_gen, "process_transformer_output"):
            @torch.no_grad()
            def personaplex_step_with_hidden(
                self_gen,
                input_tokens: torch.Tensor = None,
                moshi_tokens: torch.Tensor = None,
                text_token: torch.Tensor = None,
                depformer_replace_tokens: torch.Tensor | None = None,
            ):
                prepared = self_gen.prepare_step_input(input_tokens, moshi_tokens, text_token)
                if prepared is None:
                    return None
                input_, provided_, target_, model_input_position, target_position = prepared
                state = self_gen._streaming_state
                transformer_out, text_logits = state.graphed_main(input_)
                output = self_gen.process_transformer_output(
                    transformer_out,
                    text_logits,
                    provided_,
                    target_,
                    model_input_position,
                    target_position,
                )
                return output, transformer_out, transformer_out

            lm_gen._step = types.MethodType(personaplex_step_with_hidden, lm_gen)
            lm_gen.streaming_forever(1)
            self._warmup_runtime()
            print("[liveTryPlasticity] installed PersonaPlex graphed hidden capture", flush=True)
            return

        from moshi.models.lm import scatter_with_mask_
        from moshi.modules.transformer import create_sin_embedding
        from moshi.utils.sampling import sample_token

        capture_layer = int(self.tf_capture_layer) % len(lm_model.transformer.layers)

        old_state = getattr(lm_gen, "_streaming_state", None)
        if old_state is not None:
            with contextlib.suppress(Exception):
                old_state.__exit__(None, None, None)
            with contextlib.suppress(Exception):
                lm_gen._stop_streaming()

        def forward_text_with_layer(self_lm, sequence, sum_condition=None, cross_attention_src=None):
            B, K, S = sequence.shape
            assert K == self_lm.num_codebooks, (K, self_lm.num_codebooks)
            input_sequence = sequence
            input_ = None
            for cb_index in range(self_lm.num_audio_codebooks):
                audio_emb = self_lm.emb[cb_index](input_sequence[:, cb_index + self_lm.audio_offset])
                input_ = audio_emb if input_ is None else input_ + audio_emb
            text_emb = self_lm.text_emb(input_sequence[:, 0])
            input_ = text_emb if input_ is None else input_ + text_emb
            if sum_condition is not None:
                input_ = input_ + sum_condition.to(input_)
            if cross_attention_src is not None:
                cross_attention_src = cross_attention_src.to(input_)

            transformer = self_lm.transformer
            _, T, C = input_.shape
            dtype_input = input_.dtype
            state = transformer._streaming_state
            if state is None:
                offsets = torch.zeros(1, dtype=torch.long, device=input_.device)
            else:
                offsets = state.offsets

            x = input_
            if transformer.positional_embedding in {"sin", "sin_rope"}:
                positions = torch.arange(T, device=x.device).view(1, -1, 1)
                positions = positions + offsets.view(-1, 1, 1)
                pos_emb = create_sin_embedding(positions, C, max_period=transformer.max_period, dtype=x.dtype)
                x = x + transformer.positional_scale * pos_emb

            captured = x
            for idx, layer in enumerate(transformer.layers):
                x = layer(x, cross_attention_src=cross_attention_src)
                if idx == capture_layer:
                    captured = x

            if state is not None:
                state.offsets[:] = torch.where(state.exec_mask, state.offsets + T, state.offsets)

            transformer_out = x.to(dtype_input)
            layer_hidden = captured.to(dtype_input)
            if self_lm.out_norm:
                transformer_out = self_lm.out_norm(transformer_out)
            text_logits = self_lm.text_linear(transformer_out)
            text_logits = text_logits[:, None]
            return transformer_out, text_logits, layer_hidden

        @torch.no_grad()
        def step_with_layer(self_gen, input_tokens: torch.Tensor, depformer_replace_tokens: torch.Tensor | None = None):
            state = self_gen._streaming_state
            if state is None:
                raise RuntimeError("You should wrap those calls with a `with lm_gen.streaming(): ...`.")
            lm_model_local = self_gen.lm_model

            assert input_tokens.dim() == 3, "Shape should be [B, K, T]."
            B, Ki, S = input_tokens.shape
            assert B == state.batch_size, f"Got a batch size {B}, expected {state.batch_size}"
            assert S == 1, "Only support being given steps one by one."
            needed_tokens = lm_model_local.num_codebooks - lm_model_local.dep_q - 1
            assert Ki >= needed_tokens, f"We expect {needed_tokens} tokens from the user stream, got {Ki}."
            if Ki > needed_tokens:
                input_tokens = input_tokens[:, :needed_tokens, :]

            CT = state.cache.shape[2]
            delays = self_gen.delays_cuda[lm_model_local.dep_q + 1:]
            write_positions = (state.offsets[:, None, None] + delays[:, None]) % CT
            scatter_with_mask_(state.cache[:, lm_model_local.dep_q + 1:], -1, write_positions, input_tokens, state.exec_mask[:, None, None])

            is_init = state.offsets[:, None, None] <= self_gen.delays_cuda[:, None]
            is_init |= ~state.exec_mask[:, None, None]
            positions = (state.offsets % CT)[:, None, None].expand_as(is_init)
            input_ = state.cache.gather(dim=2, index=positions)
            input_ = torch.where(is_init, state.initial, input_)

            if self_gen.check:
                assert not (input_ == lm_model_local.ungenerated_token_id).any(), (state.offsets, input_)
                assert (input_[:, lm_model_local.audio_offset:] <= lm_model_local.card).all(), input_
                assert (input_[:, :1] <= lm_model_local.text_card).all()

            zero = torch.full((1,), lm_model_local.zero_token_id, dtype=torch.long, device=input_.device)
            if self_gen.cfg_coef != 1.:
                if state.cfg_is_masked_until is not None:
                    limit = self_gen.delays_cuda[:, None] + state.cfg_is_masked_until.view(-1, 1, 1)
                    is_zeroed = state.offsets[:, None, None] <= limit
                    masked = torch.where(is_zeroed & ~is_init, zero, input_)
                    input_ = torch.cat([input_, masked], dim=0)
                else:
                    input_ = input_.repeat(2, 1, 1)
                if self_gen.cfg_is_no_text:
                    input_[B:, :1] = torch.where(~is_init[:, :1], zero, input_[B:, :1])

            transformer_out, text_logits, layer_hidden = state.graphed_main(input_, state.condition_sum, state.condition_cross)
            if self_gen.cfg_coef != 1.:
                logits, logits_null = text_logits.chunk(2)
                if self_gen.cfg_is_no_text:
                    text_logits = logits
                    layer_hidden = layer_hidden[:B]
                else:
                    text_logits = logits_null + (logits - logits_null) * self_gen.cfg_coef
                    layer_hidden = layer_hidden[:B]

            if self_gen.on_text_logits_hook:
                self_gen.on_text_logits_hook(text_logits)
            text_token = sample_token(text_logits.float(), self_gen.use_sampling, self_gen.temp_text, self_gen.top_k_text)
            assert text_token.dim() == 3, text_token.shape
            assert text_token.shape[2] == 1
            assert text_token.shape[1] == 1, "Only one text stream supported."
            text_token = text_token[:, 0, 0]
            if self_gen.on_text_hook is not None:
                self_gen.on_text_hook(text_token)

            if state.graphed_depth is None:
                audio_tokens = None
            else:
                if depformer_replace_tokens is None:
                    audio_tokens = state.graphed_depth(text_token, transformer_out)
                else:
                    assert depformer_replace_tokens.dim() == 3
                    audio_tokens = depformer_replace_tokens.squeeze(-1)
                if self_gen.on_audio_hook is not None:
                    self_gen.on_audio_hook(audio_tokens)

            state.offsets = torch.where(state.exec_mask, state.offsets + 1, state.offsets)
            state.offset_cpu += 1
            positions = (state.offsets % CT)[:, None, None]
            scatter_with_mask_(state.cache[:, :1], -1, positions, text_token[:, None, None], state.exec_mask[:, None, None])
            if audio_tokens is not None:
                audio_tokens = audio_tokens[:, :, None]
                scatter_with_mask_(state.cache[:, 1: lm_model_local.dep_q + 1, :], -1, positions.expand_as(audio_tokens), audio_tokens, state.exec_mask[:, None, None])

            if not self_gen.support_out_of_sync and state.offset_cpu <= self_gen.max_delay:
                return None
            gen_delays_cuda = self_gen.delays_cuda[: lm_model_local.dep_q + 1]
            index = (state.offsets[:, None, None] - self_gen.max_delay + gen_delays_cuda[:, None]) % CT
            out = state.cache.gather(dim=2, index=index)
            mask = (state.offsets <= self_gen.max_delay) | ~state.exec_mask
            out[mask, :, :] = lm_model_local.ungenerated_token_id
            return out, transformer_out, layer_hidden

        lm_model.forward_text = types.MethodType(forward_text_with_layer, lm_model)
        lm_gen._step = types.MethodType(step_with_layer, lm_gen)
        lm_gen.streaming_forever(1)
        self._warmup_runtime()
        print(f"[liveTryPlasticity] installed graphed layer capture layer={self.tf_capture_layer}", flush=True)

    @torch.no_grad()
    def _step(self, pcm24: np.ndarray) -> dict:
        self.step += 1
        t0 = time.perf_counter()
        chunk = torch.from_numpy(pcm24).to(self.device, dtype=torch.float32)[None, None]

        t_encode0 = time.perf_counter()
        codes = self.mimi.encode(chunk)
        t_encode1 = time.perf_counter()
        if self.skip_first:
            self.mimi.reset_streaming()
            self.skip_first = False

        t_lm0 = time.perf_counter()
        lm_out = self.lm_gen._step(codes[:, :, :1])
        t_lm1 = time.perf_counter()

        tokens = None
        helium_hidden = None
        if lm_out is not None:
            if not (isinstance(lm_out, tuple) and len(lm_out) == 3):
                raise RuntimeError(f"Moshi graph layer[-2] contract failure: got {type(lm_out)} len={len(lm_out) if isinstance(lm_out, tuple) else 'n/a'}")
            tokens, _transformer_out, layer_hidden = lm_out
            helium_hidden = layer_hidden[:1, -1:].detach().float().cpu()

        token = -1
        token_piece = ""
        decode_ms = 0.0
        reply_codes = None
        if tokens is None:
            reply_pcm = np.zeros(MIMI_FRAME_SIZE, dtype=np.float32)
        else:
            token = int(tokens[0, 0, 0].detach().item())
            token_piece = self.decode_piece(token)
            if token_piece:
                self.audio_text += token_piece
            reply_codes = tokens[:, 1:].detach().to(device="cpu", dtype=torch.int16)
            t_decode0 = time.perf_counter()
            reply = self.mimi.decode(tokens[:, 1:])
            reply_pcm = reply[0, 0].detach().float().cpu().numpy()
            decode_ms = 1000.0 * (time.perf_counter() - t_decode0)
            if reply_pcm.shape[0] < MIMI_FRAME_SIZE:
                reply_pcm = np.pad(reply_pcm, (0, MIMI_FRAME_SIZE - reply_pcm.shape[0]))
            elif reply_pcm.shape[0] > MIMI_FRAME_SIZE:
                reply_pcm = reply_pcm[:MIMI_FRAME_SIZE]

        reply_rms = float(np.sqrt(np.mean(np.square(reply_pcm, dtype=np.float32))))
        reply_peak = float(np.max(np.abs(reply_pcm))) if reply_pcm.size else 0.0
        input_rms = float(np.sqrt(np.mean(np.square(pcm24, dtype=np.float32))))
        encode_ms = 1000.0 * (t_encode1 - t_encode0)
        lm_ms = 1000.0 * (t_lm1 - t_lm0)
        total_ms = 1000.0 * (time.perf_counter() - t0)

        reply_i16 = np.clip(reply_pcm, -1.0, 1.0)
        reply_i16 = (reply_i16 * 32767.0).astype(np.int16)
        audio_b64 = base64.b64encode(reply_i16.tobytes()).decode("ascii")

        print(
            "[liveTryStudio] moshi "
            f"step={self.step} token={token} piece={token_piece!r} "
            f"in_rms={input_rms:.5f} reply_rms={reply_rms:.5f} peak={reply_peak:.3f} "
            f"hidden={helium_hidden is not None} "
            f"encode={encode_ms:.1f}ms lm={lm_ms:.1f}ms decode={decode_ms:.1f}ms total={total_ms:.1f}ms",
            flush=True,
        )

        return {
            "step": int(self.step),
            "sample_rate": TARGET_SR,
            "reply_i16_b64": audio_b64,
            "reply_rms": reply_rms,
            "reply_peak": reply_peak,
            "input_rms": input_rms,
            "token": token,
            "piece": token_piece,
            "sampled_text": self.sampled_text,
            "audio_text": self.audio_text,
            "encode_ms": encode_ms,
            "lm_ms": lm_ms,
            "decode_ms": decode_ms,
            "total_ms": total_ms,
            "helium_hidden": helium_hidden,
            "reply_codes": reply_codes,
        }

    @torch.no_grad()
    def process_ready_steps_limited(self, max_steps: int) -> list[dict]:
        """Process a bounded number of Mimi frames.

        The base MoshiOnlyEngine drains the entire input buffer before returning.
        In live mode that is dangerous: mic audio can accumulate while Moshi is
        loading, then avatar frames do not reach the sender until the backlog is
        fully processed. Bounded draining keeps the producer/sender interleaved.
        """
        events: list[dict] = []
        for _ in range(max(1, int(max_steps))):
            if self.input_buffer.shape[0] < MIMI_FRAME_SIZE:
                break
            pcm = self.input_buffer[:MIMI_FRAME_SIZE].copy()
            self.input_buffer = self.input_buffer[MIMI_FRAME_SIZE:].copy()
            events.append(self._step(pcm))
        return events


# ---------------------------------------------------------------------------
# Utilities
# ---------------------------------------------------------------------------

def _sync_cuda() -> None:
    if torch.cuda.is_available():
        torch.cuda.synchronize()


def _ms(t0: float) -> float:
    return 1000.0 * (time.perf_counter() - t0)


def encode_jpeg_b64(frame_rgb: np.ndarray, quality: int) -> str:
    frame_bgr = cv2.cvtColor(frame_rgb, cv2.COLOR_RGB2BGR)
    ok, enc = cv2.imencode(".jpg", frame_bgr, [int(cv2.IMWRITE_JPEG_QUALITY), int(quality)])
    if not ok:
        raise RuntimeError("JPEG encode failed")
    return base64.b64encode(enc.tobytes()).decode("ascii")


def encode_jpeg_bytes(frame_rgb: np.ndarray, quality: int) -> bytes:
    frame_bgr = cv2.cvtColor(frame_rgb, cv2.COLOR_RGB2BGR)
    ok, enc = cv2.imencode(".jpg", frame_bgr, [int(cv2.IMWRITE_JPEG_QUALITY), int(quality)])
    if not ok:
        raise RuntimeError("JPEG encode failed")
    return enc.tobytes()


def _pcm_f32_to_i16_b64(pcm: np.ndarray) -> str:
    arr = np.clip(np.asarray(pcm, dtype=np.float32), -1.0, 1.0)
    return base64.b64encode((arr * 32767.0).astype(np.int16).tobytes()).decode("ascii")


def _pcm_f32_to_i16_bytes(pcm: np.ndarray) -> bytes:
    arr = np.clip(np.asarray(pcm, dtype=np.float32), -1.0, 1.0)
    return (arr * 32767.0).astype(np.int16).tobytes()


def split_audio_into_frame_slices(pcm: np.ndarray, fps: float) -> list[np.ndarray]:
    frame_samples = int(round(TARGET_SR / float(fps)))
    arr = np.asarray(pcm, dtype=np.float32)
    n_frames = max(0, int(round(arr.shape[0] / frame_samples)))
    if n_frames == 0:
        return []
    total = n_frames * frame_samples
    if arr.shape[0] < total:
        arr = np.pad(arr, (0, total - arr.shape[0]))
    elif arr.shape[0] > total:
        arr = arr[:total]
    return [arr[i * frame_samples:(i + 1) * frame_samples].copy() for i in range(n_frames)]


def load_audio_24k(path: str) -> np.ndarray:
    # torchaudio 2.11 has no built-in decoders (needs torchcodec) -> use soundfile,
    # with an ffmpeg fallback for containers soundfile cannot read.
    import soundfile as sf
    try:
        wav, sr = sf.read(path, dtype="float32", always_2d=True)
    except Exception:
        import subprocess, tempfile
        with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as f:
            tmp = f.name
        subprocess.run(["ffmpeg", "-y", "-v", "error", "-i", path, "-ac", "1",
                        "-ar", str(TARGET_SR), "-f", "wav", tmp], check=True)
        wav, sr = sf.read(tmp, dtype="float32", always_2d=True)
        os.remove(tmp)
    wav = wav.mean(axis=1)
    if sr != TARGET_SR:
        import librosa
        wav = librosa.resample(wav, orig_sr=sr, target_sr=TARGET_SR)
    return np.asarray(wav, dtype=np.float32)


def load_ref_image(path: str | Path, device: torch.device, dtype: torch.dtype) -> torch.Tensor:
    img = Image.open(path).convert("RGB").resize((512, 512), Image.LANCZOS)
    return T.ToTensor()(img).unsqueeze(0).to(device=device, dtype=dtype)


# ---------------------------------------------------------------------------
# FM + Renderer weight loading (identical to liveTryFM.py)
# ---------------------------------------------------------------------------

def _clean_generator_state(ckpt: dict) -> dict:
    raw = ckpt.get("ema_state_dict") or ckpt.get("state_dict", ckpt.get("model", ckpt))
    if isinstance(raw, dict) and "model" in raw and isinstance(raw["model"], dict):
        raw = raw["model"]
    return {k.replace("model.", "", 1) if k.startswith("model.") else k: v for k, v in raw.items()}


def _load_fm(args: argparse.Namespace, device: torch.device) -> FMGenerator:
    t_total = time.perf_counter()
    fm = FMGenerator(args).to(device).eval()
    ckpt = torch.load(args.generator_path, map_location="cpu")
    cleaned = _clean_generator_state(ckpt)
    missing, unexpected = fm.load_state_dict(cleaned, strict=False)
    print(
        f"[liveTryHeliumFM][FM] base loaded missing={len(missing)} unexpected={len(unexpected)}",
        flush=True,
    )
    lora_path = str(getattr(args, "lora_generator_path", "") or "")
    if lora_path:
        apply_lora_to_model(
            fm,
            rank=int(getattr(args, "lora_rank", 64) or 64),
            alpha=float(getattr(args, "lora_alpha", 128) or 128),
            dropout=float(getattr(args, "lora_dropout", 0.05)),
            include_pose_lora=not bool(getattr(args, "no_lora_pose_projection", False)),
            include_audio_lora=not bool(getattr(args, "no_lora_audio_projection", False)),
            only_pose_lora=bool(getattr(args, "only_lora_pose_projection", False)),
        )
        lora_ckpt = torch.load(lora_path, map_location="cpu")
        lora_cleaned = _clean_generator_state(lora_ckpt)
        missing_lora, unexpected_lora = fm.load_state_dict(lora_cleaned, strict=False)
        lora_keys = sum(1 for key in lora_cleaned if "lora_" in key)
        print(
            f"[liveTryHeliumFM][FM] lora loaded path={lora_path} "
            f"lora_keys={lora_keys} missing={len(missing_lora)} unexpected={len(unexpected_lora)}",
            flush=True,
        )
    fm.to(device).eval()
    _sync_cuda()
    print(f"[liveTryHeliumFM][FM] loaded in {_ms(t_total):.0f}ms", flush=True)
    return fm


def _load_renderer(args: argparse.Namespace, device: torch.device, dtype: torch.dtype) -> IMTRenderer:
    t_total = time.perf_counter()
    renderer = IMTRenderer(args).to(device).eval()
    ckpt = torch.load(args.renderer_path, map_location="cpu")
    raw = ckpt.get("state_dict", ckpt.get("model", ckpt))
    cleaned = {k.replace("gen.", "", 1).replace("model.", "", 1): v for k, v in raw.items()}
    missing, unexpected = renderer.load_state_dict(cleaned, strict=False)
    renderer = renderer.to(dtype=dtype)
    _sync_cuda()
    if getattr(args, "compile_renderer", False):
        @torch.no_grad()
        def _fused_render(motion_latent, g_r, m_r, f_r):
            ta_c = renderer.adapt(motion_latent, g_r)
            m_c = renderer.latent_token_decoder(ta_c)
            frames = renderer.decode(m_c, m_r, f_r)
            return frames
        renderer._fused_render = torch.compile(_fused_render)
    print(
        f"[liveTryHeliumFM][renderer] loaded in {_ms(t_total):.0f}ms "
        f"missing={len(missing)} unexpected={len(unexpected)}",
        flush=True,
    )
    return renderer


# ---------------------------------------------------------------------------
# Helium extractor (chunk-local, batch LM)
# ---------------------------------------------------------------------------

class HeliumExtractor:
    """Prefix-growing raw Helium + global interpolation.

    This matches the best Stage 3 diagnostic:
      audio prefix -> raw Helium -> append only new raw steps
      -> one global interpolation over accumulated raw Helium
      -> emit last target_frames
    """

    def __init__(
        self,
        helium_mimi: "MimiModel",
        helium_lm: "LMModel",
        device: torch.device,
    ) -> None:
        self.helium_mimi = helium_mimi
        self.helium_lm = helium_lm
        self.device = device
        self._lock = threading.Lock()
        self._prefix_pcm = np.empty(0, dtype=np.float32)
        self._raw_parts: list[torch.Tensor] = []
        self._prev_raw_len = 0
        self._emitted_frames = 0

    def reset(self) -> None:
        with self._lock:
            self._prefix_pcm = np.empty(0, dtype=np.float32)
            self._raw_parts = []
            self._prev_raw_len = 0
            self._emitted_frames = 0

    def _extract_raw(self, pcm_np: np.ndarray) -> torch.Tensor:
        wav = torch.from_numpy(np.asarray(pcm_np, dtype=np.float32)).to(self.device, dtype=torch.float32)[None, None]

        codes = self.helium_mimi.encode(wav)
        codes = codes[:, :MAIN_CODEBOOKS, :].detach()
        batch_size, n_q, total_steps = codes.shape

        dtype = next(self.helium_lm.parameters()).dtype
        input_emb = torch.zeros(
            batch_size, total_steps, self.helium_lm.dim, device=self.device, dtype=dtype
        )
        for q in range(n_q):
            input_emb = input_emb + self.helium_lm.emb[q](codes[:, q].long())

        padding_ids = torch.full(
            (batch_size, total_steps),
            self.helium_lm.existing_text_padding_id,
            dtype=torch.long,
            device=self.device,
        )
        input_emb = input_emb + self.helium_lm.text_emb(padding_ids)

        if getattr(self.helium_lm.transformer, "_streaming_state", None) is not None:
            raise RuntimeError("helium_lm must stay in batch mode (non-streaming)")

        captured: list[torch.Tensor] = []

        def _hook(_mod, _inp, out):
            captured.append(out.detach())

        handle = self.helium_lm.transformer.layers[-2].register_forward_hook(_hook)
        try:
            self.helium_lm.transformer(input_emb)
        finally:
            handle.remove()

        if len(captured) != 1:
            raise RuntimeError(f"Helium hook captured {len(captured)} tensors; expected 1")
        return captured[0].squeeze(0).float().contiguous()  # [T_raw, 4096]

    @torch.no_grad()
    def extract_raw_chunk(self, pcm_np: np.ndarray) -> torch.Tensor:
        """Return only the new raw 12.5Hz Helium steps for one new audio chunk."""
        pcm = np.asarray(pcm_np, dtype=np.float32)
        if pcm.ndim != 1 or pcm.size == 0:
            raise RuntimeError("HeliumExtractor.extract_raw_chunk expects non-empty 1D PCM")

        with self._lock:
            self._prefix_pcm = np.concatenate([self._prefix_pcm, pcm], axis=0)
            raw_prefix = self._extract_raw(self._prefix_pcm)
            new_raw = raw_prefix[self._prev_raw_len:]
            if int(new_raw.shape[0]) == 0 and int(raw_prefix.shape[0]) > 0:
                new_raw = raw_prefix[-1:]
            self._raw_parts.append(new_raw.cpu())
            self._prev_raw_len = int(raw_prefix.shape[0])
            return new_raw.contiguous()

    @torch.no_grad()
    def extract_exact_chunk_from_prefix(
        self,
        pcm_prefix: np.ndarray,
        chunk_start_frame: int,
        target_frames: int,
    ) -> torch.Tensor:
        """Return the raw 12.5Hz Helium slice for a video-frame window.

        This path is used by file-mode lookahead. It must return raw Helium
        steps so the studio adapter remains the only temporal upsampler.
        """
        pcm = np.asarray(pcm_prefix, dtype=np.float32)
        if pcm.ndim != 1 or pcm.size == 0:
            raise RuntimeError("HeliumExtractor.extract_exact_chunk_from_prefix expects non-empty 1D PCM")
        with self._lock:
            raw_prefix = self._extract_raw(pcm)
        start_frame = int(chunk_start_frame)
        end_frame = start_frame + int(target_frames)
        start_raw = int(round(start_frame * 0.5))
        end_raw = int(round(end_frame * 0.5))
        start_raw = max(0, min(start_raw, int(raw_prefix.shape[0])))
        end_raw = max(start_raw + 1, min(end_raw, int(raw_prefix.shape[0])))
        if end_raw > int(raw_prefix.shape[0]):
            raise RuntimeError(
                f"Requested raw slice [{start_raw}, {end_raw}) exceeds prefix Helium length {raw_prefix.shape[0]}"
            )
        return raw_prefix[start_raw:end_raw].contiguous()


# ---------------------------------------------------------------------------
# Main engine
# ---------------------------------------------------------------------------

class LiveHeliumFMEngine:
    """Helium extraction + FM + renderer, session-stateful."""

    def __init__(self, args: argparse.Namespace) -> None:
        if args.tf32:
            torch.backends.cuda.matmul.allow_tf32 = True
            torch.backends.cudnn.allow_tf32 = True
            torch.set_float32_matmul_precision("high")

        self.args = args
        self.device = torch.device(args.device)
        renderer_precision = str(
            getattr(args, "renderer_precision", "fp32")
        ).lower()
        self.dtype = {
            "fp32": torch.float32,
            "fp16": torch.float16,
            "bf16": torch.bfloat16,
        }[renderer_precision]
        self.fps = float(args.fps)
        self.audio_chunk_sec = float(getattr(args, "audio_chunk_sec", 0.96))
        self.audio_chunk_samples = int(round(self.audio_chunk_sec * TARGET_SR))
        self.fm_chunk_frames = max(1, int(getattr(args, "fm_chunk_frames", 24)))
        self.live_sliding_window = bool(getattr(args, "enable_live_sliding_window", False))
        self.slide_past_frames = max(0, int(getattr(args, "slide_past_frames", 10)))
        self.slide_future_frames = max(0, int(getattr(args, "slide_future_frames", 3)))
        self.render_sub_batch = max(1, int(args.render_sub_batch))
        self.jpeg_quality = int(args.jpeg_quality)
        # emotion conditioning (audio + semantic, 80 ms streaming state)
        self.emo = LiveEmotionState(args)
        self.e_cfg_scale = float(getattr(args, "e_cfg_scale", 3.0))
        trained_window = int(round(float(args.wav2vec_sec) * self.fps))
        if self.fm_chunk_frames != trained_window:
            print(
                f"[liveTryHeliumFM] WARNING fm_chunk_frames={self.fm_chunk_frames} "
                f"but wav2vec_sec*fps={trained_window}",
                flush=True,
            )
        if self.live_sliding_window:
            print(
                f"[liveTryHeliumFM][typeAC] IMTalker lookahead enabled "
                f"past={self.slide_past_frames}f current={self.fm_chunk_frames}f "
                f"future={self.slide_future_frames}f",
                flush=True,
            )

        t_total = time.perf_counter()

        # FM + renderer
        self.fm = _load_fm(args, self.device)
        self.renderer = _load_renderer(args, self.device, self.dtype)

        # Adapter: either the six-layer projected-frontend model or UniTalk's
        # 12-layer model with only its final layer passed to IMTalker.
        t_adapter = time.perf_counter()
        if bool(int(getattr(args, "audio_direct", 0))):
            self.studio_adapter = None
            print("[liveTryHeliumFM] audio_direct: Helium->wav2vec adapter skipped",
                  flush=True)
        elif args.adapter_type == "unitalk_last_layer":
            self.studio_adapter = UniTalkLastLayerLiveAdapter(
                args.wav2vec_model_path,
                args.adapter_dropout,
            ).to(self.device).float().eval()
        else:
            self.studio_adapter = StudioNativeLiveAdapter(
                args.wav2vec_model_path,
                args.adapter_num_layers,
                args.adapter_dropout,
            ).to(self.device).float().eval()
        if self.studio_adapter is not None:
            payload = torch.load(args.adapter_path, map_location="cpu")
            if isinstance(payload, dict) and args.adapter_type == "frontend":
                saved_args = payload.get("args", {})
                if saved_args and int(saved_args.get("num_layers", args.adapter_num_layers)) != int(args.adapter_num_layers):
                    print(
                        f"[liveTryHeliumFrontendFM] WARNING checkpoint num_layers={saved_args.get('num_layers')} "
                        f"but CLI adapter_num_layers={args.adapter_num_layers}",
                        flush=True,
                    )
                state = payload.get("adapter", payload.get("model", payload))
            else:
                state = payload
            self.studio_adapter.load_state_dict(state, strict=True)
            _sync_cuda()
            print(
                f"[liveTryHeliumFrontendFM][adapter] type={args.adapter_type} loaded in {_ms(t_adapter):.0f}ms "
                f"path={args.adapter_path}",
                flush=True,
            )

        # Raw HF Wav2Vec2 target path used during Helium adapter training.
        t_w2v = time.perf_counter()
        self.wav2vec_feature_extractor = Wav2Vec2FeatureExtractor.from_pretrained(
            args.wav2vec_model_path,
            local_files_only=True,
        )
        self.wav2vec_model = self.studio_adapter.wav2vec if self.studio_adapter is not None else None
        _sync_cuda()
        print(
            f"[liveTryHeliumStudioFM][wav2vec] loaded in {_ms(t_w2v):.0f}ms "
            f"path={args.wav2vec_model_path}",
            flush=True,
        )

        # Reference image: pre-compute identity + motion-ref features once
        ref_tensor = load_ref_image(args.ref_path, self.device, self.dtype)
        with torch.no_grad():
            self.f_r, self.g_r = self.renderer.dense_feature_encoder(ref_tensor)
            self.ref_x = self.renderer.latent_token_encoder(ref_tensor).to(dtype=torch.float32)
            ta_r = self.renderer.adapt(self.ref_x.to(dtype=self.dtype), self.g_r)
            self.m_r = self.renderer.latent_token_decoder(ta_r)
        _sync_cuda()
        self.eye_blink_enabled = bool(getattr(args, "enable_eye_blink_composite", False))
        self._blink_maps: tuple[torch.Tensor, ...] | None = None
        self._eye_masks: tuple[torch.Tensor, ...] | None = None
        self._render_frame_cursor: int = 0
        if self.eye_blink_enabled:
            self._init_eye_blink_composite()

        # Moshi models for Helium extraction
        self._init_moshi(args)

        # Optional local audio file (simulate-live mode)
        self.audio_pcm: np.ndarray | None = None
        if getattr(args, "audio_path", "") and Path(args.audio_path).is_file():
            self.audio_pcm = load_audio_24k(args.audio_path)
            print(
                f"[liveTryHeliumFM] audio_path loaded: {self.audio_pcm.shape[0]/TARGET_SR:.2f}s "
                f"chunk={self.audio_chunk_sec:.3f}s/{self.audio_chunk_samples} samples "
                f"fm_chunk={self.fm_chunk_frames}f",
                flush=True,
            )

        # Shared noise tensor (pre-generated, indexed by absolute frame position)
        self.noise_buf: torch.Tensor | None = None
        if getattr(args, "shared_noise", False):
            max_frames = int(getattr(args, "noise_max_frames", 5000))
            gen = torch.Generator(device=self.device)
            gen.manual_seed(int(getattr(args, "noise_seed", 1234)))
            self.noise_buf = torch.randn(
                1, max_frames, int(args.dim_w), device=self.device, generator=gen
            )
            print(f"[liveTryHeliumFM] shared noise buf: {tuple(self.noise_buf.shape)}", flush=True)

        # Per-session state (reset on each new client)
        self.stream_state: dict | None = None
        self.abs_frame: int = 0
        self.helium_context_tail: torch.Tensor | None = None
        self.helium_deque_size: int = 100
        self.helium_deque: torch.Tensor | None = None
        self.helium_deque_filled: int = 0
        self.silence_helium_seed: torch.Tensor | None = None
        silence_helium_path = str(
            getattr(args, "silence_helium_path", "") or ""
        )
        if silence_helium_path:
            payload = torch.load(silence_helium_path, map_location="cpu")
            seed = (
                payload.get("silence_helium_mean")
                if isinstance(payload, dict)
                else payload
            )
            if not isinstance(seed, torch.Tensor) or seed.numel() != 4096:
                raise RuntimeError(
                    f"Invalid silence Helium seed: {silence_helium_path}"
                )
            self.silence_helium_seed = seed.reshape(1, 4096).to(
                device=self.device, dtype=torch.float32
            )
            print(
                f"[liveTryHeliumFM] silence Helium seed loaded: "
                f"{silence_helium_path}",
                flush=True,
            )
        self._pcm_accum: np.ndarray = np.empty(0, dtype=np.float32)
        self.dump_motion = bool(getattr(args, "dump_motion", False))
        self.dump_dir = Path(getattr(args, "dump_dir", ROOT / "live_try_dumps"))
        self._session_motion_parts: list[torch.Tensor] = []
        self._session_helium_parts: list[torch.Tensor] = []
        self._session_adapter_50_parts: list[torch.Tensor] = []
        self._session_adapter_25_parts: list[torch.Tensor] = []
        self._session_projected_audio_parts: list[torch.Tensor] = []
        self._session_audio_parts: list[np.ndarray] = []
        self._session_live_token_parts: list[torch.Tensor] = []
        self._session_chunk_rows: list[dict] = []
        self._session_reply_events: list[dict] = []
        self._session_started_wall: float = time.time()

        # JPEG encoding thread pool (CPU-only work, parallelizable)
        self._jpeg_pool = concurrent.futures.ThreadPoolExecutor(
            max_workers=4, thread_name_prefix="jpeg"
        )

        # Warmup
        self._warmup()

        print(
            f"[liveTryHeliumFM] ready — total startup {_ms(t_total):.0f}ms "
            f"fm_chunk={self.fm_chunk_frames} render_sub={self.render_sub_batch} "
            f"dtype={self.dtype}",
            flush=True,
        )

    def _init_moshi(self, args: argparse.Namespace) -> None:
        if bool(int(getattr(args, "audio_direct", 0))):
            self.extractor = None
            print("[liveTryHeliumFM] audio_direct: Helium extractor skipped "
                  "(reply PCM -> IMTalker wav2vec)", flush=True)
            return
        if bool(getattr(args, "direct_reply_hidden", False)) and bool(getattr(args, "enable_moshi_reply", False)):
            self.extractor = None
            print("[liveTryHeliumStudioFM] using direct Moshi reply hidden; batch Helium extractor skipped", flush=True)
            return

        from generate_helium import load_mimi_and_lm

        t0 = time.perf_counter()
        helium_mimi, helium_lm, _ = load_mimi_and_lm(args)
        helium_mimi.eval()
        helium_lm.eval()
        print(
            f"[liveTryHeliumFM] Moshi loaded in {_ms(t0):.0f}ms "
            f"dtype={next(helium_lm.parameters()).dtype}",
            flush=True,
        )

        self.extractor = HeliumExtractor(helium_mimi, helium_lm, self.device)

    def reset_session(self) -> None:
        """Call when a new WebSocket client connects or sends 'start'."""
        self.stream_state = None
        self.abs_frame = 0
        self._render_frame_cursor = 0
        if getattr(self, "emo", None) is not None:
            self.emo.reset()
        self.helium_context_tail = None
        self.helium_deque = None
        self.helium_deque_filled = 0
        self._pcm_accum = np.empty(0, dtype=np.float32)
        if self.extractor is not None:
            self.extractor.reset()
        self._session_motion_parts = []
        self._session_helium_parts = []
        self._session_adapter_50_parts = []
        self._session_adapter_25_parts = []
        self._session_projected_audio_parts = []
        self._session_audio_parts = []
        self._session_live_token_parts = []
        self._session_chunk_rows = []
        self._session_reply_events = []
        self._session_started_wall = time.time()

    @torch.no_grad()
    def _warmup(self) -> None:
        dummy_pcm = np.zeros(self.audio_chunk_samples, dtype=np.float32)

        t0 = time.perf_counter()
        if self.extractor is None:
            raw_steps = max(1, int(round(self.fm_chunk_frames * 12.5 / float(self.fps))))
            if self.silence_helium_seed is not None:
                dummy_helium = self.silence_helium_seed.expand(
                    raw_steps, -1
                ).contiguous()
            else:
                dummy_helium = torch.zeros(
                    raw_steps, 4096, device=self.device, dtype=torch.float32
                )
            print(f"[liveTryHeliumStudioFM][warmup] raw_helium=skipped direct_hidden raw_steps={raw_steps}", flush=True)
        else:
            dummy_helium = self.extractor.extract_raw_chunk(dummy_pcm)
            _sync_cuda()
            print(f"[liveTryHeliumStudioFM][warmup] raw_helium={_ms(t0):.0f}ms", flush=True)
            self.extractor.reset()

        t0 = time.perf_counter()
        motion, _info = self._sample_motion_from_helium(dummy_helium, self.fm_chunk_frames)
        _sync_cuda()
        print(f"[liveTryHeliumFM][warmup] fm={_ms(t0):.0f}ms motion={tuple(motion.shape)}", flush=True)
        self.stream_state = None
        self.abs_frame = 0
        self._render_frame_cursor = 0
        self.helium_context_tail = None
        self.helium_deque = None
        self.helium_deque_filled = 0

        t0 = time.perf_counter()
        dummy_motion = torch.zeros(self.render_sub_batch, 32, device=self.device, dtype=self.dtype)
        _frames, _timings = self._render_motion(dummy_motion)
        _sync_cuda()
        print(f"[liveTryHeliumFM][warmup] renderer={_ms(t0):.0f}ms", flush=True)
        self._render_frame_cursor = 0

        # Warmup JPEG pool
        t0 = time.perf_counter()
        dummy_np = np.zeros((512, 512, 3), dtype=np.uint8)
        _ = encode_jpeg_b64(dummy_np, self.jpeg_quality)
        print(f"[liveTryHeliumFM][warmup] jpeg={_ms(t0):.0f}ms", flush=True)

        self.stream_state = None

    def feed_pcm(self, pcm_s16le_bytes: bytes) -> Optional[tuple[torch.Tensor, dict, np.ndarray]]:
        pcm = np.frombuffer(pcm_s16le_bytes, dtype=np.int16).astype(np.float32) / 32768.0
        return self.feed_pcm_f32(pcm)

    def feed_pcm_f32(self, pcm_f32: np.ndarray) -> Optional[tuple[torch.Tensor, dict, np.ndarray]]:
        pcm = np.asarray(pcm_f32, dtype=np.float32)
        self._pcm_accum = np.concatenate([self._pcm_accum, pcm])
        if self._pcm_accum.shape[0] < self.audio_chunk_samples:
            return None
        chunk = self._pcm_accum[:self.audio_chunk_samples].copy()
        self._pcm_accum = self._pcm_accum[self.audio_chunk_samples:]
        motion, info = self._process_pcm_chunk(chunk, self.fm_chunk_frames)
        self._record_session_chunk(chunk, motion, info)
        return motion, info, chunk

    @torch.no_grad()
    def _process_pcm_chunk(self, pcm_chunk: np.ndarray, target_frames: int) -> tuple[torch.Tensor, dict]:
        timings: dict = {}
        target_frames = max(1, min(int(target_frames), self.fm_chunk_frames))
        pcm_chunk = np.asarray(pcm_chunk, dtype=np.float32)

        # ---- SILENCE must be STILL (closed mouth) -------------------------
        # All-zero audio is out-of-distribution for wav2vec (it normalises per
        # call), and the generator answered digital silence with a WIDE-OPEN
        # mouth. The drain lead-in and every pause therefore showed the avatar
        # "talking" with no audio. Render the reference latent instead: the
        # head stays still and the mouth stays shut.
        rms = float(np.sqrt(np.mean(np.square(pcm_chunk)))) if pcm_chunk.size else 0.0
        if bool(int(getattr(self.args, "idle_still", 1))) and rms < float(
                getattr(self.args, "idle_rms_threshold", 0.004)):
            self.pending_audio = pcm_chunk
            idle = self.ref_x.reshape(1, -1).to(self.device, dtype=torch.float32)
            motion = idle.expand(int(target_frames), -1).clone()
            tail = getattr(self, "_last_motion_tail", None)
            cf = int(getattr(self.args, "idle_crossfade_frames", 10))
            if tail is not None and cf > 0 and tail.shape[0] > 0:
                k = min(cf, int(target_frames), int(tail.shape[0]))
                w = torch.linspace(0.0, 1.0, k + 2, device=motion.device,
                                   dtype=motion.dtype)[1:-1].view(-1, 1)
                motion[:k] = (1.0 - w) * tail[-k:].to(motion.device, motion.dtype) + w * motion[:k]
            self._last_motion_tail = motion[-min(10, int(target_frames)):].detach().clone()
            self.abs_frame += int(target_frames)
            timings.update({
                "frames": int(target_frames),
                "abs_start": int(self.abs_frame) - int(target_frames),
                "helium_ms": 0.0,
                "fm_ms": 0.0,
                "idle": 1,
            })
            return motion, timings

        t0 = time.perf_counter()
        if self.extractor is None:
            helium = None
            timings["helium_ms"] = 0.0
        else:
            helium = self.extractor.extract_raw_chunk(pcm_chunk)
            timings["helium_ms"] = _ms(t0)
        self.pending_audio = np.asarray(pcm_chunk, dtype=np.float32)
        motion, fm_info = self._sample_motion_from_helium(helium, target_frames)
        self._last_motion_tail = motion[-min(10, int(motion.shape[0])):].detach().clone()
        timings.update(fm_info)
        return motion, timings

    def _emo_conditioning(self, n_frames: int) -> dict:
        """Per-frame emotion conditioning for the next chunk.

        All of the schedule logic lives in LiveEmotionState, which reproduces
        the per-frame schedule used by the offline conversation renders.
        """
        emo = getattr(self, "emo", None)
        if emo is None or not getattr(emo, "enable", False):
            return {}
        return emo.conditioning(int(n_frames)) or {}

    @torch.no_grad()
    def _sample_motion_from_helium(self, helium: torch.Tensor, target_frames: int) -> tuple[torch.Tensor, dict]:
        timings: dict = {}

        # ---- direct-audio path (our IMTalker checkpoint): reply PCM -> wav2vec
        # -> generator. Skips Helium + the Helium->wav2vec adapter entirely.
        if bool(getattr(self.args, "audio_direct", False)):
            import torchaudio as _ta
            pending = getattr(self, "pending_audio", None)
            if pending is None:
                pending = np.zeros(int(self.audio_chunk_samples), dtype=np.float32)
            a24 = torch.from_numpy(np.asarray(pending, dtype=np.float32)).view(1, -1)
            a16 = _ta.functional.resample(a24, 24000, 16000)  # (1, T) for wav2vec
            _prev_tail = getattr(self, "_prev_a_tail", None)
            # 3 s of LEFT CONTEXT for the wav2vec encoder (see FM.sample):
            # features from isolated chunks do not match the training features.
            _hist = getattr(self, "_fm_audio_hist", None)
            if _hist is not None and _hist.numel():
                _ctx = torch.cat([_hist, a16], dim=-1)
            else:
                _ctx = a16
            if _ctx.shape[-1] > 16000 * 5:
                _ctx = _ctx[:, -16000 * 5:]
            # Encode the left context (better wav2vec features) but feed the
            # generator ONLY the frames of THIS chunk: the generator is trained
            # on 50-frame (2 s) clips. Passing the full 5 s context made it
            # generate 125 frames, and the caller kept the FIRST 50 -> motion
            # paired with audio ~3 s in the past ("mouth does not match").
            _tc = int(round(_ctx.shape[-1] * 25.0 / 16000.0))
            _feat = self.fm.audio_encoder.inference(
                _ctx.to(self.device), seq_len=_tc)
            if _feat.ndim == 3:
                _feat = _feat[0]
            _tf = int(target_frames)
            _feat = _feat[-_tf:] if _feat.shape[0] > _tf else _feat
            if _feat.shape[0] < _tf:
                _feat = torch.cat(
                    [_feat, _feat[-1:].expand(_tf - _feat.shape[0], -1)], dim=0)
            data: dict = {"a": a16, "a_feat": _feat[None], "ref_x": self.ref_x,
                          "gen_frames": _tf}
            if _prev_tail is not None:
                data["prev_a_wav"] = _prev_tail
            try:
                self._prev_a_tail = a16[:, -6400:].detach().clone()
                self._fm_audio_hist = _ctx[:, -(16000 * 3):].detach().clone()
            except Exception:
                self._prev_a_tail = None
                self._fm_audio_hist = None
            emo_cond = self._emo_conditioning(int(target_frames))
            if emo_cond:
                data.update(emo_cond)
                timings["emo_cls"] = int(emo_cond["_cls"])
                timings["emo_inten"] = round(float(emo_cond["_inten"]), 3)
                emo_cond.pop("_cls", None)
                emo_cond.pop("_inten", None)
            t0 = time.perf_counter()
            _call = dict(
                a_cfg_scale=float(self.args.a_cfg_scale),
                e_cfg_scale=float(getattr(self, "e_cfg_scale", 3.0)),
                nfe=int(self.args.nfe),
            )
            try:
                motion, self.stream_state = self.fm.sample(
                    data, stream_state=self.stream_state, return_state=True, **_call
                )
            except RuntimeError as _e:
                # The frozen transformer's prev-context path has a shape bug for
                # some chunk sizes; fall back to an independent chunk (same as the
                # offline render) instead of dropping the whole stream.
                print(f"[audio_direct] state carry failed ({_e}); retry cold", flush=True)
                self.stream_state = None
                motion = self.fm.sample(data, **_call)
            timings["fm_ms"] = _ms(t0)
            timings["audio_direct"] = 1
            motion = motion.squeeze(0).detach()
            if motion.shape[0] > target_frames:
                motion = motion[-target_frames:]
            else:
                motion = motion[:target_frames]
            if int(motion.shape[0]) < target_frames:
                pad = motion[-1:].expand(target_frames - motion.shape[0], -1)
                motion = torch.cat([motion, pad], 0)
            timings["frames"] = int(motion.shape[0])
            timings["abs_start"] = self.abs_frame
            self.abs_frame += timings["frames"]
            timings["helium_ms"] = 0.0
            return motion, timings

        t_adapter = time.perf_counter()
        helium = helium.to(self.device, dtype=torch.float32).contiguous()
        target_frames = int(target_frames)
        current_steps = int(helium.shape[0])
        deque_size = int(getattr(self, "helium_deque_size", 100))
        if self.helium_deque is None:
            if self.silence_helium_seed is not None:
                self.helium_deque = self.silence_helium_seed.expand(
                    deque_size, -1
                ).clone()
            else:
                self.helium_deque = torch.zeros(
                    deque_size,
                    helium.shape[1],
                    device=self.device,
                    dtype=torch.float32,
                )
            self.helium_deque_filled = 0
        if current_steps >= deque_size:
            self.helium_deque = helium[-deque_size:].detach().clone()
            self.helium_deque_filled = deque_size
        else:
            self.helium_deque = torch.cat([self.helium_deque[current_steps:], helium], dim=0).contiguous()
            self.helium_deque_filled = min(deque_size, int(self.helium_deque_filled) + current_steps)

        adapter_window_mode = str(
            getattr(self.args, "adapter_window_mode", "tail")
        )
        if adapter_window_mode == "lookahead":
            # Match training: process the full 8-second window at 50 Hz,
            # emit .96 seconds, and retain .48 seconds as future context.
            future_steps = int(getattr(self.args, "adapter_future_steps", 6))
            target_len_50_full = deque_size * 4
            _baseline, _cnn, feat_50_full = self.studio_adapter.forward_single(
                self.helium_deque, target_len_50_full
            )
            if self.live_sliding_window:
                # Type AC: keep the Type A Helium deque + adapter path, but
                # give IMTalker/FM a small [past + current + future] feature
                # window and emit only current frames. This places lookahead in
                # IMTalker instead of discarding future context inside adapter.
                past_25 = int(self.slide_past_frames)
                future_25 = int(self.slide_future_frames)
                current_25 = int(target_frames)
                past_50 = past_25 * 2
                future_50 = future_25 * 2
                current_50 = current_25 * 2
                full_len_50 = int(feat_50_full.shape[0])
                current_end_50 = full_len_50 - future_50
                current_start_50 = current_end_50 - current_50
                window_start_50 = current_start_50 - past_50
                window_end_50 = current_end_50 + future_50
                if window_start_50 < 0 or current_start_50 < 0 or window_end_50 > full_len_50:
                    raise RuntimeError(
                        "Sliding adapter window is out of range: "
                        f"full={full_len_50} start={window_start_50} "
                        f"current_start={current_start_50} end={window_end_50}"
                    )
                feat_50 = feat_50_full[window_start_50:window_end_50].contiguous()
                feat_25 = F.interpolate(
                    feat_50.T.unsqueeze(0),
                    size=past_25 + current_25 + future_25,
                    mode="linear",
                    align_corners=False,
                ).squeeze(0).T.contiguous()
            else:
                emitted_frames_50 = max(1, current_steps * 4)
                future_frames_50 = max(0, future_steps * 4)
                segment_end = int(feat_50_full.shape[0]) - future_frames_50
                segment_start = segment_end - emitted_frames_50
                if segment_start < 0:
                    raise RuntimeError(
                        "Look-ahead adapter output is shorter than its emit/future region"
                    )
                feat_50 = feat_50_full[segment_start:segment_end].contiguous()
                feat_25 = F.interpolate(
                    feat_50.T.unsqueeze(0),
                    size=target_frames,
                    mode="linear",
                    align_corners=False,
                ).squeeze(0).T.contiguous()
        else:
            # Legacy behavior: predict at 25 Hz and emit the newest tail.
            target_len_25_full = deque_size * 2
            _baseline, _cnn, feat_25_full = self.studio_adapter.forward_single(
                self.helium_deque, target_len_25_full
            )
            fresh_frames = max(1, current_steps * 2)
            if int(feat_25_full.shape[0]) < fresh_frames:
                raise RuntimeError(
                    f"Deque adapter output too short: got "
                    f"{feat_25_full.shape[0]}, need {fresh_frames}"
                )
            feat_25 = feat_25_full[-fresh_frames:].contiguous()
            feat_50 = feat_25
        if (not self.live_sliding_window) and int(feat_25.shape[0]) != target_frames:
            feat_25 = F.interpolate(
                feat_25.T.unsqueeze(0),
                size=target_frames,
                mode="linear",
                align_corners=False,
            ).squeeze(0).T.contiguous()
        projected_a = self.fm._project_audio(feat_25.unsqueeze(0).float())
        timings["adapter_ms"] = _ms(t_adapter)
        timings["helium_ms"] = timings["adapter_ms"]
        timings["helium_deque_filled"] = int(self.helium_deque_filled)

        data: dict = {"a_feat": feat_25.unsqueeze(0).float(), "ref_x": self.ref_x}
        emo_cond = self._emo_conditioning(int(target_frames))
        if emo_cond:
            data.update(emo_cond)
            timings["emo_cls"] = int(emo_cond["_cls"])
            timings["emo_inten"] = round(float(emo_cond["_inten"]), 3)
            emo_cond.pop("_cls", None)
            emo_cond.pop("_inten", None)
        if self.noise_buf is not None:
            if self.live_sliding_window:
                noise_start = max(0, self.abs_frame - int(self.slide_past_frames))
                noise_end = noise_start + int(feat_25.shape[0])
            else:
                noise_start = self.abs_frame
                noise_end = self.abs_frame + target_frames
            data["noise_init"] = self.noise_buf[:, noise_start:noise_end]
        t_fm = time.perf_counter()
        if self.live_sliding_window:
            motion = self.fm.sample(
                data,
                a_cfg_scale=float(self.args.a_cfg_scale),
                e_cfg_scale=float(getattr(self, "e_cfg_scale", 3.0)),
                nfe=int(self.args.nfe),
            )
        else:
            motion, self.stream_state = self.fm.sample(
                data,
                a_cfg_scale=float(self.args.a_cfg_scale),
                e_cfg_scale=float(getattr(self, "e_cfg_scale", 3.0)),
                nfe=int(self.args.nfe),
                stream_state=self.stream_state,
                return_state=True,
            )
        timings["fm_ms"] = _ms(t_fm)

        motion = motion.squeeze(0).detach()
        if self.live_sliding_window:
            start = int(self.slide_past_frames)
            motion = motion[start:start + target_frames]
        else:
            motion = motion[:target_frames]
        ref_blend = float(getattr(self.args, "motion_ref_blend", 0.0) or 0.0)
        if ref_blend > 0.0:
            ref_motion = self.ref_x.detach().float()
            if ref_motion.ndim == 3:
                ref_motion = ref_motion[0]
            if ref_motion.ndim == 1:
                ref_motion = ref_motion.unsqueeze(0)
            if int(ref_motion.shape[0]) == 1:
                ref_motion = ref_motion.expand(int(motion.shape[0]), -1)
            elif int(ref_motion.shape[0]) != int(motion.shape[0]):
                ref_motion = F.interpolate(
                    ref_motion.T.unsqueeze(0),
                    size=int(motion.shape[0]),
                    mode="linear",
                    align_corners=False,
                ).squeeze(0).T
            ref_motion = ref_motion.to(device=motion.device, dtype=motion.dtype)
            blend = max(0.0, min(1.0, ref_blend))
            motion = motion.mul(1.0 - blend).add(ref_motion, alpha=blend)
        timings["helium_feat"] = helium.detach().cpu()
        timings["adapter_feat_50"] = feat_50.detach().cpu()
        timings["adapter_feat_25"] = feat_25.detach().cpu()
        timings["projected_audio"] = projected_a.squeeze(0).detach().cpu()
        timings["frames"] = int(motion.shape[0])
        timings["abs_start"] = self.abs_frame
        self.abs_frame += timings["frames"]
        return motion, timings

    def _record_session_chunk(self, pcm_chunk: np.ndarray, motion: torch.Tensor, info: dict) -> None:
        self._session_audio_parts.append(np.asarray(pcm_chunk, dtype=np.float32).copy())
        self._session_motion_parts.append(motion.detach().float().cpu().clone())
        helium_feat = info.get("helium_feat")
        if isinstance(helium_feat, torch.Tensor):
            self._session_helium_parts.append(helium_feat.float().cpu().clone())
        adapter_feat_50 = info.get("adapter_feat_50")
        if isinstance(adapter_feat_50, torch.Tensor):
            self._session_adapter_50_parts.append(adapter_feat_50.float().cpu().clone())
        adapter_feat_25 = info.get("adapter_feat_25")
        if isinstance(adapter_feat_25, torch.Tensor):
            self._session_adapter_25_parts.append(adapter_feat_25.float().cpu().clone())
        projected_audio = info.get("projected_audio")
        if isinstance(projected_audio, torch.Tensor):
            self._session_projected_audio_parts.append(projected_audio.float().cpu().clone())
        self._session_chunk_rows.append({
            "chunk": len(self._session_chunk_rows) + 1,
            "abs_start": int(info.get("abs_start", 0)),
            "frames": int(info.get("frames", int(motion.shape[0]))),
            "samples": int(len(pcm_chunk)),
            "helium_ms": float(info.get("helium_ms", 0.0)),
            "fm_ms": float(info.get("fm_ms", 0.0)),
        })

    @torch.no_grad()
    def _extract_wav2vec_raw_50hz(self, audio_24k: np.ndarray) -> torch.Tensor:
        arr = np.asarray(audio_24k, dtype=np.float32)
        if arr.ndim != 1 or arr.size == 0:
            return torch.empty((0, 768), dtype=torch.float32)
        wav = torch.from_numpy(arr).view(1, -1)
        wav16 = torchaudio.functional.resample(wav, TARGET_SR, WAV2VEC_SR).squeeze(0).contiguous().numpy()
        inputs = self.wav2vec_feature_extractor(
            wav16,
            sampling_rate=WAV2VEC_SR,
            return_tensors="pt",
            padding=True,
        )
        kwargs = {
            "input_values": inputs.input_values.to(self.device),
        }
        if getattr(inputs, "attention_mask", None) is not None:
            kwargs["attention_mask"] = inputs.attention_mask.to(self.device)
        frontend = self.wav2vec_model.extract_projected_frontend(**kwargs)
        feat = self.wav2vec_model.encode_from_projected_frontend(
            frontend
        ).last_hidden_state.detach().float().cpu()[0].contiguous()
        return feat

    def dump_last_session(self, *, source: str = "") -> Optional[Path]:
        if not self.dump_motion or not self._session_motion_parts:
            return None
        self.dump_dir.mkdir(parents=True, exist_ok=True)
        session_dir = self.dump_dir / "last_session"
        session_dir.mkdir(parents=True, exist_ok=True)

        motion = torch.cat(self._session_motion_parts, dim=0).contiguous()
        audio = np.concatenate(self._session_audio_parts, axis=0) if self._session_audio_parts else np.empty(0, dtype=np.float32)

        motion_path = session_dir / "full_motion.pt"
        helium_path = session_dir / "full_helium_raw.pt"
        adapter_50_path = session_dir / "full_adapter_w2v_50hz.pt"
        adapter_25_path = session_dir / "full_adapter_w2v_25fps.pt"
        wav2vec_50_path = session_dir / "full_wav2vec_50hz.pt"
        projected_audio_path = session_dir / "full_projected_audio_32.pt"
        audio_path = session_dir / "full_moshi_reply_24k.wav"
        live_tokens_path = session_dir / "live_mimi_tokens.pt"
        reply_events_path = session_dir / "reply_events.jsonl"
        reply_text_path = session_dir / "reply_text.txt"
        meta_path = session_dir / "meta.json"
        helium = None
        adapter_50 = None
        adapter_25 = None
        wav2vec_50 = None
        projected_audio = None
        live_tokens = None
        if self._session_helium_parts:
            helium = torch.cat(self._session_helium_parts, dim=0).contiguous()
        if self._session_adapter_50_parts:
            adapter_50 = torch.cat(self._session_adapter_50_parts, dim=0).contiguous()
        if self._session_adapter_25_parts:
            adapter_25 = torch.cat(self._session_adapter_25_parts, dim=0).contiguous()
        if self._session_projected_audio_parts:
            projected_audio = torch.cat(self._session_projected_audio_parts, dim=0).contiguous()
        if self._session_live_token_parts:
            live_tokens = torch.cat(self._session_live_token_parts, dim=2).contiguous()
        if audio.size > 0:
            wav2vec_50 = self._extract_wav2vec_raw_50hz(audio)

        torch.save({
            "motion": motion,
            "chunks": self._session_chunk_rows,
            "fps": float(self.fps),
            "audio_chunk_sec": float(self.audio_chunk_sec),
            "fm_chunk_frames": int(self.fm_chunk_frames),
            "audio_feat_dim": int(getattr(self.args, "audio_feat_dim", 768)),
            "audio_adapter_dim": int(getattr(self.args, "audio_adapter_dim", 512)),
            "wav2vec_sec": float(self.args.wav2vec_sec),
            "ref_path": str(self.args.ref_path),
            "generator_path": str(self.args.generator_path),
            "renderer_path": str(self.args.renderer_path),
            "source": source,
        }, motion_path)
        if helium is not None:
            torch.save({
                "helium": helium,
                "chunks": self._session_chunk_rows,
                "fps": float(self.fps),
                "audio_chunk_sec": float(self.audio_chunk_sec),
                "fm_chunk_frames": int(self.fm_chunk_frames),
                "audio_feat_dim": int(getattr(self.args, "audio_feat_dim", 4096)),
                "source": source,
            }, helium_path)
        if adapter_50 is not None:
            torch.save({
                "adapter_feat_50": adapter_50,
                "chunks": self._session_chunk_rows,
                "source": source,
            }, adapter_50_path)
        if adapter_25 is not None:
            torch.save({
                "adapter_feat_25": adapter_25,
                "chunks": self._session_chunk_rows,
                "fps": float(self.fps),
                "source": source,
            }, adapter_25_path)
        if wav2vec_50 is not None:
            torch.save({
                "wav2vec_50hz": wav2vec_50,
                "chunks": self._session_chunk_rows,
                "sample_rate": int(WAV2VEC_SR),
                "source": source,
            }, wav2vec_50_path)
        if projected_audio is not None:
            torch.save({
                "projected_audio": projected_audio,
                "chunks": self._session_chunk_rows,
                "fps": float(self.fps),
                "source": source,
            }, projected_audio_path)
        if live_tokens is not None:
            torch.save({
                "live_mimi_tokens": live_tokens,
                "chunks": self._session_chunk_rows,
                "source": source,
            }, live_tokens_path)
        if audio.size > 0:
            torchaudio.save(str(audio_path), torch.from_numpy(audio).view(1, -1), TARGET_SR)
        if self._session_reply_events:
            with reply_events_path.open("w", encoding="utf-8") as f:
                for row in self._session_reply_events:
                    f.write(json.dumps(row, ensure_ascii=False) + "\n")
            reply_text = "".join(str(row.get("piece", "")) for row in self._session_reply_events)
            reply_text_path.write_text(reply_text, encoding="utf-8")
        meta_path.write_text(json.dumps({
            "source": source,
            "session_started_wall": float(self._session_started_wall),
            "fps": float(self.fps),
            "audio_chunk_sec": float(self.audio_chunk_sec),
            "fm_chunk_frames": int(self.fm_chunk_frames),
            "motion_frames": int(motion.shape[0]),
            "helium_frames": int(helium.shape[0]) if helium is not None else 0,
            "adapter_50_frames": int(adapter_50.shape[0]) if adapter_50 is not None else 0,
            "adapter_25_frames": int(adapter_25.shape[0]) if adapter_25 is not None else 0,
            "wav2vec_50_frames": int(wav2vec_50.shape[0]) if wav2vec_50 is not None else 0,
            "projected_audio_frames": int(projected_audio.shape[0]) if projected_audio is not None else 0,
            "audio_samples": int(audio.shape[0]),
            "audio_seconds": float(audio.shape[0] / TARGET_SR) if audio.size > 0 else 0.0,
            "reply_text_chars": int(sum(len(str(row.get("piece", ""))) for row in self._session_reply_events)),
            "reply_events": int(len(self._session_reply_events)),
            "chunks": self._session_chunk_rows,
            "ref_path": str(self.args.ref_path),
            "generator_path": str(self.args.generator_path),
            "renderer_path": str(self.args.renderer_path),
        }, indent=2), encoding="utf-8")
        print(f"[liveTryHeliumFM] dumped last session -> {session_dir}", flush=True)
        return session_dir

    def _extract_motion_tensor_from_payload(self, payload, path: str) -> torch.Tensor:
        if isinstance(payload, torch.Tensor):
            motion = payload
        elif isinstance(payload, dict):
            candidates = ["motion", "motion_latents", "latents", "full_motion", "pred_motion", "x"]
            motion = None
            for key in candidates:
                value = payload.get(key)
                if isinstance(value, torch.Tensor):
                    motion = value
                    break
            if motion is None:
                tensor_items = [
                    value
                    for value in payload.values()
                    if isinstance(value, torch.Tensor) and value.ndim >= 2 and value.shape[-1] == int(self.args.dim_w)
                ]
                if tensor_items:
                    motion = tensor_items[0]
            if motion is None:
                raise ValueError(
                    f"No motion tensor found in blink_motion_path={path}. "
                    f"Available keys={list(payload.keys())[:20]}"
                )
        else:
            raise TypeError(f"Unsupported blink motion payload type {type(payload)} from {path}")

        motion = motion.detach().float().cpu()
        if motion.ndim == 3 and motion.shape[0] == 1:
            motion = motion[0]
        if motion.ndim != 2 or motion.shape[-1] != int(self.args.dim_w):
            raise ValueError(
                f"Blink motion must have shape (T,{int(self.args.dim_w)}) or (1,T,{int(self.args.dim_w)}); "
                f"got {tuple(motion.shape)} from {path}"
            )
        return motion.contiguous()

    def _make_eye_mask(self, height: int, width: int) -> torch.Tensor:
        yy = torch.linspace(0.0, 1.0, int(height), device=self.device, dtype=self.dtype).view(1, 1, height, 1)
        xx = torch.linspace(0.0, 1.0, int(width), device=self.device, dtype=self.dtype).view(1, 1, 1, width)
        center_y = float(getattr(self.args, "eye_center_y", 0.405))
        radius_x = max(float(getattr(self.args, "eye_radius_x", 0.145)), 1e-6)
        radius_y = max(float(getattr(self.args, "eye_radius_y", 0.070)), 1e-6)
        feather = max(float(getattr(self.args, "eye_feather", 0.10)), 1e-6)

        mask = torch.zeros(1, 1, height, width, device=self.device, dtype=self.dtype)
        for center_x in (
            float(getattr(self.args, "eye_left_x", 0.36)),
            float(getattr(self.args, "eye_right_x", 0.64)),
        ):
            dist = ((xx - center_x) / radius_x).square() + ((yy - center_y) / radius_y).square()
            ellipse = ((1.0 + feather) - dist).clamp(0.0, feather) / feather
            mask = torch.maximum(mask, ellipse)
        return mask.clamp(0.0, 1.0).contiguous()

    @torch.no_grad()
    def _init_eye_blink_composite(self) -> None:
        blink_path = str(getattr(self.args, "blink_motion_path", "") or "")
        if not blink_path:
            raise ValueError("--enable_eye_blink_composite requires --blink_motion_path")
        if not Path(blink_path).is_file():
            raise FileNotFoundError(f"blink_motion_path does not exist: {blink_path}")

        payload = torch.load(blink_path, map_location="cpu")
        blink_motion = self._extract_motion_tensor_from_payload(payload, blink_path)
        blink_maps_parts: list[list[torch.Tensor]] = []
        chunk = max(1, int(getattr(self.args, "render_sub_batch", 8)))
        for start in range(0, int(blink_motion.shape[0]), chunk):
            sub = blink_motion[start:start + chunk].to(self.device, dtype=self.dtype)
            g_sub = self.g_r.expand(int(sub.shape[0]), -1)
            ta_b = self.renderer.adapt(sub, g_sub)
            maps_b = self.renderer.latent_token_decoder(ta_b)
            if not blink_maps_parts:
                blink_maps_parts = [[] for _ in range(len(maps_b))]
            for idx, map_b in enumerate(maps_b):
                blink_maps_parts[idx].append(map_b.detach())

        self._blink_maps = tuple(torch.cat(parts, dim=0).contiguous() for parts in blink_maps_parts)
        self._eye_masks = tuple(self._make_eye_mask(m.shape[-2], m.shape[-1]) for m in self._blink_maps)
        _sync_cuda()
        shapes = [tuple(m.shape) for m in self._blink_maps]
        print(
            f"[liveTryHeliumFM][blink] cached blink maps from {blink_path}: "
            f"frames={int(blink_motion.shape[0])} shapes={shapes}",
            flush=True,
        )

    def _composite_eye_blink_maps(
        self,
        current_maps: tuple[torch.Tensor, ...] | list[torch.Tensor],
        start_frame: int,
        num_frames: int,
    ) -> tuple[torch.Tensor, ...]:
        if self._blink_maps is None or self._eye_masks is None:
            return tuple(current_maps)
        blink_len = int(self._blink_maps[0].shape[0])
        if blink_len <= 0:
            return tuple(current_maps)
        indices = (torch.arange(int(num_frames), device=self.device) + int(start_frame)) % blink_len
        composited: list[torch.Tensor] = []
        for cur, blink_all, mask in zip(current_maps, self._blink_maps, self._eye_masks):
            blink = blink_all.index_select(0, indices).to(device=cur.device, dtype=cur.dtype)
            mask = mask.to(device=cur.device, dtype=cur.dtype)
            composited.append(blink * mask + cur * (1.0 - mask))
        return tuple(composited)

    @torch.no_grad()
    def _render_motion(self, motion: torch.Tensor) -> tuple[np.ndarray, dict]:
        timings: dict = {}
        t_total = time.perf_counter()
        motion = motion.to(self.device, dtype=self.dtype)
        n = int(motion.shape[0])

        g_r_sub = self.g_r.expand(n, -1)
        m_r_sub = tuple(m.expand(n, -1, -1, -1) for m in self.m_r)
        f_r_sub = [f.expand(n, -1, -1, -1) for f in self.f_r]

        render_start = self._render_frame_cursor
        fused = getattr(self.renderer, '_fused_render', None)
        if self.eye_blink_enabled:
            fused = None
        if fused is not None:
            frames = fused(motion, g_r_sub, m_r_sub, f_r_sub)
        else:
            ta_c = self.renderer.adapt(motion, g_r_sub)
            m_c = self.renderer.latent_token_decoder(ta_c)
            if self.eye_blink_enabled:
                m_c = self._composite_eye_blink_maps(m_c, render_start, n)
            frames = self.renderer.decode(m_c, m_r_sub, f_r_sub)
        self._render_frame_cursor += n

        frames_np = frames.detach().float().clamp(0, 1).mul(255).to(torch.uint8)
        frames_np = frames_np.permute(0, 2, 3, 1).contiguous().cpu().numpy()
        timings["total_ms"] = _ms(t_total)
        return frames_np, timings

    def render_and_encode_subbatch(
        self,
        motion_sub: torch.Tensor,
        audio_slices: list[np.ndarray],
        abs_start: int,
        text_payload: str,
        avatar_chunk_id: int,
        total_gen_ms: float,
    ) -> list[dict]:
        """Render a sub-batch of frames, JPEG-encode in parallel, return packet dicts."""
        frames_np, _render_info = self._render_motion(motion_sub)

        jpeg_futures = []
        for frame_rgb in frames_np:
            jpeg_futures.append(
                self._jpeg_pool.submit(encode_jpeg_bytes, frame_rgb, self.jpeg_quality)
            )

        packets = []
        gen_ms_i = int(round(float(total_gen_ms)))
        sr_i = int(round(float(TARGET_SR)))
        for j, fut in enumerate(jpeg_futures):
            idx = abs_start + j
            audio_slice = audio_slices[j] if j < len(audio_slices) else np.zeros(
                int(round(TARGET_SR / self.fps)), dtype=np.float32
            )
            jpeg_bytes = fut.result()
            output_audio_codec = str(
                getattr(self.args, "output_audio_codec", "pcm")
            ).lower()
            pcm_b = (
                b""
                if output_audio_codec == "opus"
                else _pcm_f32_to_i16_bytes(audio_slice)
            )
            if os.environ.get("SESSION_DUMP") == "1" and idx < 150:
                try:
                    _dir = "/workspace/session_frames"
                    os.makedirs(_dir, exist_ok=True)
                    with open(f"{_dir}/f{idx:04d}.jpg", "wb") as _f:
                        _f.write(jpeg_bytes)
                except Exception:
                    pass
            blob = _wsbin.pack_av_frame(
                idx,
                idx + 1,
                gen_ms_i,
                sr_i,
                jpeg_bytes,
                pcm_b,
                text_payload,
                int(avatar_chunk_id),
            )
            packets.append(
                {
                    "frame_number": idx,
                    "ws_kind": "bytes",
                    "data": blob,
                    "audio_pcm": np.asarray(audio_slice, dtype=np.float32).copy(),
                    "t_ready": time.perf_counter(),
                }
            )
        return packets

    def audio_slice(self, frame_idx: int) -> np.ndarray:
        frame_samples = int(round(TARGET_SR / self.fps))
        src = self.audio_pcm
        if src is None and self._session_audio_parts:
            # live reply mode: the audio is recorded per chunk instead
            src = np.concatenate(self._session_audio_parts, axis=0)
        if src is None:
            return np.zeros(frame_samples, dtype=np.float32)
        start = frame_idx * frame_samples
        chunk = src[start:start + frame_samples]
        if chunk.shape[0] < frame_samples:
            chunk = np.pad(chunk, (0, frame_samples - chunk.shape[0]))
        return chunk


# ---------------------------------------------------------------------------
# WebSocket streaming coroutine (file-driven mode, unchanged)
# ---------------------------------------------------------------------------

async def stream_from_file(ws: WebSocket, engine: LiveHeliumFMEngine) -> None:
    """Simulate live streaming using --audio_path, sending frames back over WS."""
    if engine.audio_pcm is None:
        await ws.send_json({"type": "error", "msg": "No --audio_path given for file-streaming mode"})
        return

    audio = engine.audio_pcm
    lookahead_chunks = max(0, int(getattr(engine.args, "file_chunk_lookahead", 0)))
    total_chunks = int(np.ceil(len(audio) / engine.audio_chunk_samples))
    start_wall = time.perf_counter()
    emitted = 0

    print(
        f"[liveTryHeliumFM] stream_from_file: {total_chunks} chunks "
        f"lookahead={lookahead_chunks}",
        flush=True,
    )

    async def _emit_motion_chunk(
        motion: torch.Tensor,
        fm_info: dict,
        chunk_label: int,
        emitted_so_far: int,
    ) -> int:
        helium_ms = float(fm_info["helium_ms"])
        fm_ms = float(fm_info["fm_ms"])
        n_frames = int(motion.shape[0])
        all_frames_np: list[np.ndarray] = []
        render_ms = 0.0
        for sb_start in range(0, n_frames, engine.render_sub_batch):
            sub = motion[sb_start:sb_start + engine.render_sub_batch].to(
                engine.device, dtype=engine.dtype
            )
            frames_np, render_info = engine._render_motion(sub)
            render_ms += float(render_info["total_ms"])
            all_frames_np.extend(frames_np)

        print(
            f"[liveTryHeliumFM][chunk#{chunk_label}] "
            f"helium={helium_ms:.0f}ms fm={fm_ms:.0f}ms "
            f"render={render_ms:.0f}ms frames={n_frames} "
            f"emo={fm_info.get('emo_cls', '-')}/{fm_info.get('emo_inten', '-')} "
            f"abs_start={fm_info['abs_start']}",
            flush=True,
        )

        for j, frame_rgb in enumerate(all_frames_np):
            idx = emitted_so_far + j
            chunk_id = idx + 1
            audio_b64 = _pcm_f32_to_i16_b64(engine.audio_slice(idx))
            jpeg_b64 = encode_jpeg_b64(frame_rgb, engine.jpeg_quality)

            if str(getattr(self.args, "output_audio_codec", "pcm")).lower() != "pcm":
                # with pcm the frames carry their own audio -> sending it again
                # as JSON made the client schedule both, inserting silence into
                # the playback cursor and pushing the audio ahead of the mouth
                await ws.send_json({
                    "type": "chunk_audio",
                    "chunk_id": chunk_id,
                    "sample_rate": TARGET_SR,
                    "pcm_s16le_b64": audio_b64,
                })
            if os.environ.get("SESSION_DUMP") == "1" and chunk_id <= 120:
                try:
                    import base64 as _b64
                    _dir = "/workspace/session_frames"
                    os.makedirs(_dir, exist_ok=True)
                    with open(f"{_dir}/f{chunk_id:04d}.jpg", "wb") as _f:
                        _f.write(_b64.b64decode(jpeg_b64))
                except Exception:
                    pass
            await ws.send_json({
                "type": "chunk_frame",
                "chunk_id": chunk_id,
                "frame_idx": 0,
                "jpeg_b64": jpeg_b64,
                "moshi_text": (
                    f"Helium+FM | chunk#{chunk_label} "
                    f"helium={helium_ms:.0f}ms fm={fm_ms:.0f}ms "
                    f"render={render_ms:.0f}ms"
                ),
                "server_fps": round(float(engine.fps), 1),
                "chunks_done": chunk_label,
            })
            target_t = start_wall + (idx + 1) / engine.fps
            await asyncio.sleep(max(0.0, target_t - time.perf_counter()))
        return emitted_so_far + len(all_frames_np)

    if lookahead_chunks <= 0:
        for chunk_idx in range(total_chunks):
            pcm_chunk = audio[
                chunk_idx * engine.audio_chunk_samples:(chunk_idx + 1) * engine.audio_chunk_samples
            ]
            pcm_real = pcm_chunk.copy()
            target_frames = int(round(pcm_chunk.shape[0] * engine.fps / TARGET_SR))
            if pcm_chunk.shape[0] < engine.audio_chunk_samples:
                pcm_chunk = np.pad(pcm_chunk, (0, engine.audio_chunk_samples - pcm_chunk.shape[0]))

            motion, fm_info = engine._process_pcm_chunk(pcm_chunk, target_frames)
            engine._record_session_chunk(pcm_real, motion, fm_info)
            emitted = await _emit_motion_chunk(motion, fm_info, chunk_idx + 1, emitted)
    else:
        pending_real: list[np.ndarray] = []
        prefix_audio = np.empty(0, dtype=np.float32)
        chunk_counter = 0

        def _process_exact_pending(pcm_real: np.ndarray) -> tuple[torch.Tensor, dict]:
            nonlocal prefix_audio
            target_frames = int(round(pcm_real.shape[0] * engine.fps / TARGET_SR))
            t0 = time.perf_counter()
            helium = engine.extractor.extract_exact_chunk_from_prefix(
                prefix_audio,
                engine.abs_frame,
                target_frames,
            )
            _sync_cuda()
            fm_info: dict = {"helium_ms": _ms(t0)}
            motion, sample_info = engine._sample_motion_from_helium(helium, target_frames)
            fm_info.update(sample_info)
            return motion, fm_info

        for chunk_idx in range(total_chunks):
            pcm_real = audio[
                chunk_idx * engine.audio_chunk_samples:(chunk_idx + 1) * engine.audio_chunk_samples
            ].copy()
            prefix_audio = np.concatenate([prefix_audio, pcm_real], axis=0)
            pending_real.append(pcm_real)

            while len(pending_real) > lookahead_chunks:
                chunk_counter += 1
                oldest = pending_real.pop(0)
                motion, fm_info = _process_exact_pending(oldest)
                engine._record_session_chunk(oldest, motion, fm_info)
                emitted = await _emit_motion_chunk(motion, fm_info, chunk_counter, emitted)

        while pending_real:
            chunk_counter += 1
            oldest = pending_real.pop(0)
            motion, fm_info = _process_exact_pending(oldest)
            engine._record_session_chunk(oldest, motion, fm_info)
            emitted = await _emit_motion_chunk(motion, fm_info, chunk_counter, emitted)

    await ws.send_json({"type": "stream_end", "total_frames": emitted})
    engine.dump_last_session(source=str(engine.args.audio_path))
    print(f"[liveTryHeliumFM] stream done: {emitted} frames", flush=True)


# ---------------------------------------------------------------------------
# Options
# ---------------------------------------------------------------------------
_EMO_RULES = [
    (0, ["unacceptable", "ridiculous", "furious", "angry", "never", "stop it"]),
    (1, ["disgusting", "gross", "awful", "revolting", "sickening"]),
    (2, ["afraid", "scared", "worried", "danger", "fear", "terrified"]),
    (3, ["great", "wonderful", "glad", "happy", "love", "fantastic", "excellent", "yay", "!" ]),
    (5, ["sorry", "sad", "unfortunate", "loss", "miss", "bad news", "heartbreaking"]),
    (6, ["wow", "amazing", "surprising", "really?", "unbelievable", "oh my"]),
    (7, ["whatever", "sure you did", "as if", "right", "clearly"]),
]


def _text_emotion(text: str) -> tuple[int, float]:
    """Cheap demo-grade semantic emotion: keyword match -> (class, weight)."""
    t = (text or "").lower()
    best, score = 4, 0
    for cls, words in _EMO_RULES:
        hits = sum(1 for w in words if w in t)
        if hits > score:
            best, score = cls, hits
    return best, min(1.0, 0.45 + 0.2 * score)


_EMO_NAMES = {
    "anger": 0, "angry": 0, "disgust": 1, "fear": 2, "afraid": 2,
    "happy": 3, "joy": 3, "neutral": 4, "sad": 5, "sadness": 5,
    "surprise": 6, "surprised": 6, "contempt": 7, "shy": 7,
}


def _smoothstep(x: float) -> float:
    return x * x * (3.0 - 2.0 * x)


class LiveEmotionState:
    """Emotion conditioning for the live IMTalker generator.

    Mirrors the offline conversation demos that worked (render_conversation.py):
      * one-hot emotion per frame,
      * smoothstep blend across a switch over `blend` frames (15 = 0.6 s),
      * intensity eased in/out over `ramp` frames (12 = 0.48 s),
      * intensities around 0.3-0.5 (1.0 saturates),
      * emo_scale 1.0, e_cfg ~2.5.
    The schedule is keyed by absolute reply-audio time, so a chunk gets the
    exact per-frame slice it would have had in the single continuous render,
    instead of a chunk-constant class that pops at the boundary.
    """

    def __init__(self, args) -> None:
        self._args = args
        self.neutral_as_happy = int(getattr(args, "neutral_as_happy", 1))
        self.enable = bool(int(getattr(args, "emo_enable", 1)))
        self.mode = str(getattr(args, "emo_mode", "auto") or "auto")
        self.manual_cls = int(getattr(args, "emo_class", -1))
        self.manual_inten = float(getattr(args, "emo_inten", 0.5))
        self.gain = float(getattr(args, "emo_gain", 1.0))
        # 0.8+ is caricature; conversations look best around 0.3-0.55
        self.inten_max = float(getattr(args, "emo_inten_max", 0.55))
        self.ema = float(getattr(args, "emo_ema", 0.35))
        self.fps = float(getattr(args, "fps", 25.0))
        self.blend = max(0, int(getattr(args, "emo_blend_frames", 15)))
        self.ramp = max(0, int(getattr(args, "emo_ramp_frames", 12)))
        self.emo_scale = float(getattr(args, "emo_scale", 1.0))
        self.t = 0.0
        # autotext: emotion turns learned from the reply TEXT stream
        self.turns: list = []          # (start_sec, cls, inten)
        self.pending_text = ""
        self._last_piece_t = None
        self._last_piece_wall = None
        self.lookahead_ms = float(getattr(args, "emo_lookahead_ms", 400.0))
        self.active = True
        self.cls = 4
        self.inten = 0.0
        self.prob = np.zeros(8, np.float32)
        self.prob[4] = 1.0
        self._prev_vec = np.zeros(8, np.float32)
        self._prev_vec[4] = 1.0
        self._prev_inten = 0.0
        # script entries: (start_sec, class, intensity)
        self.script: list[tuple[float, int, float]] = []
        raw = str(getattr(args, "emo_script", "") or "")
        for part in raw.split(","):
            part = part.strip()
            if not part or "@" not in part:
                continue
            name, rest = part.split("@", 1)
            name = name.strip().lower()
            it = None
            if ":" in rest:
                rest, it_s = rest.split(":", 1)
                try:
                    it = float(it_s)
                except ValueError:
                    it = None
            try:
                t_s = float(rest)
            except ValueError:
                continue
            cls = int(name) if name.lstrip("-").isdigit() else _EMO_NAMES.get(name, 4)
            if not (0 <= cls <= 7):
                cls = 4
            self.script.append(
                (t_s, cls, float(self.manual_inten if it is None else it))
            )
        self.script.sort(key=lambda x: x[0])

    def set_manual(self, mode: str = "manual", cls: int = -1, inten: float = 0.5):
        self.mode = str(mode)
        self.manual_cls = int(cls)
        self.manual_inten = float(inten)

    def reset(self) -> None:
        self.t = 0.0
        self.active = True
        self.cls = 4
        self.inten = 0.0
        self.prob = np.zeros(8, np.float32)
        self.prob[4] = 1.0
        self._prev_vec = np.zeros(8, np.float32)
        self._prev_vec[4] = 1.0
        self._prev_inten = 0.0

    def update(self, used_audio, used_steps, assistant_active: bool) -> None:
        """Chunk-level target for auto mode (script mode ignores this)."""
        self.active = bool(assistant_active)
        if self.mode != "auto" or not self.enable:
            return
        if not self.active:
            tgt_cls, tgt_inten = 4, 0.0
        else:
            text = " ".join(
                str(s.get("audio_text") or s.get("sampled_text") or s.get("piece") or "")
                for s in (used_steps or [])
            )
            cls, w = _text_emotion(text)
            rms = (
                float(np.mean([float(s.get("reply_rms", 0.0)) for s in used_steps]))
                if used_steps else 0.0
            )
            tgt_cls = cls
            tgt_inten = min(1.0, max(0.0, (0.25 + 2.5 * rms) * w)) * self.gain
        onehot = np.zeros(8, np.float32)
        onehot[tgt_cls] = 1.0
        self.prob = (1.0 - self.ema) * self.prob + self.ema * onehot
        self.inten = (1.0 - self.ema) * self.inten + self.ema * tgt_inten
        self.cls = int(self.prob.argmax())

    def _schedule_window(self, t_start: float, n_frames: int):
        """Per-frame (one-hot, intensity) over [t_start, t_start + n/fps).

        One clean crossfade per boundary: the first `blend` frames of a segment
        move from the previous emotion to the new one, then hold; intensity
        eases from the previous segment's value over `ramp` frames. (The offline
        helper blends forward into already-blended frames, which re-triggers on
        short segments and washes the emotion out.)
        """
        if not self.script:
            return None
        fps = self.fps
        n = int(n_frames)
        g0 = int(round(t_start * fps))
        frames = g0 + np.arange(n)
        sch = np.zeros((n, 8), np.float32)
        inten = np.zeros(n, np.float32)

        starts = [int(round(s[0] * fps)) for s in self.script]
        idx = np.searchsorted(starts, frames, side="right") - 1
        idx = np.clip(idx, 0, len(self.script) - 1)

        for i in range(len(self.script)):
            sel = np.where(idx == i)[0]
            if sel.size == 0:
                continue
            off = frames[sel] - starts[i]
            p = np.clip((off + 1) / (self.blend + 1), 0.0, 1.0)
            a_b = p * p * (3.0 - 2.0 * p)
            q = np.clip((off + 1) / (self.ramp + 1), 0.0, 1.0)
            a_r = q * q * (3.0 - 2.0 * q)

            prev_cls = self.script[i - 1][1] if i > 0 else 4
            prev_it = self.script[i - 1][2] if i > 0 else 0.0
            cls, it = self.script[i][1], self.script[i][2]
            prev_oh = np.zeros(8, np.float32)
            prev_oh[prev_cls] = 1.0
            cur_oh = np.zeros(8, np.float32)
            cur_oh[cls] = 1.0
            sch[sel] = (
                (1.0 - a_b)[:, None] * prev_oh[None, :]
                + a_b[:, None] * cur_oh[None, :]
            )
            inten[sel] = (1.0 - a_r) * prev_it + a_r * it

        return sch, inten

    # ---------------- autotext: emotion from the reply text stream ----------------
    _LEX = {
        3: ("great", "good to see", "love", "awesome", "congrats", "glad", "excited",
            "perfect", "spirit", "good luck", "enjoy", "nice", "happy for"),
        6: ("really?", "wow", "what happened", "unbelievable", "seriously", "no way",
            "won't believe", "guess what"),
        0: ("worse", "damn", "annoying", "frustrat", "angry", "ridiculous",
            "moved the whole", "screwed", "fed up"),
        1: ("gross", "awful", "terrible", "disgust", "yuck", "nasty",
            "can't stand", "boring"),
        5: ("sorry", "rough", "tough", "unfortunately", "unfair", "poor you",
            "okay with that", "hang in"),
    }
    _LEX_ORDER = (6, 1, 0, 5, 3)

    # HF label -> our class index (0 Angry .. 7 Contempt)
    _HF_MAP = {
        "anger": 0, "disgust": 1, "fear": 2, "joy": 3, "neutral": 4,
        "sadness": 5, "surprise": 6, "surprised": 6, "happiness": 3,
        "sad": 5, "happy": 3, "angry": 0, "contempt": 7,
    }
    _CLF = None
    _CLF_TRIED = False
    _CLF_BUSY = False

    def _load_clf(self):
        """Load the text-emotion classifier on CPU (GPU is busy rendering)."""
        if self._CLF is not None or type(self)._CLF_TRIED:
            return
        type(self)._CLF_TRIED = True
        try:
            from transformers import pipeline
            import torch as _torch
            _torch.set_num_threads(2)
            type(self)._CLF = pipeline(
                "text-classification",
                model="j-hartmann/emotion-english-distilroberta-base",
                top_k=None, device=-1,      # CPU: never contend with the renderer
            )
            # warm-up (first inference is much slower than the rest)
            for _ in range(2):
                type(self)._CLF("warm up sentence for the emotion classifier")
            print("[emo-autotext] CPU emotion classifier loaded + warmed", flush=True)
        except Exception as e:
            print(f"[emo-autotext] classifier load failed ({e}); using lexicon", flush=True)
            type(self)._CLF = None

    def _ensure_clf_worker(self):
        if getattr(self, "_clf_worker", None) is not None:
            return
        import threading
        import queue as _queue
        self._clf_worker_q = _queue.Queue()
        self._clf_lock = threading.Lock()

        def _worker():
            self._load_clf()
            while True:
                item = self._clf_worker_q.get()
                if item is None:
                    break
                text, t_now = item
                try:
                    cls, inten = self._classify_text(text)
                    if int(getattr(self._args, "neutral_as_happy", 1)) and cls == 4:
                        cls = 3
                        inten = min(max(inten, 0.22), 0.30)
                except Exception as _e:
                    print(f"[emo-autotext] classify error: {_e}", flush=True)
                    cls, inten = 4, 0.25
                dur = min(6.0, max(1.0, len(text) * 0.055))
                start = t_now - dur
                with self._clf_lock:
                    if self.turns:
                        start = max(start, self.turns[-1][0] + 1.0)
                    if self.turns and self.turns[-1][1] == cls:
                        self.turns[-1] = (min(self.turns[-1][0], start), cls,
                                          max(self.turns[-1][2], inten))
                    else:
                        self.turns.append((start, cls, inten))
                        print(f"[emo-autotext] turn@{start:.1f}s cls={cls} "
                              f"inten={inten:.2f} | {text[:48]!r}", flush=True)

        self._clf_worker = threading.Thread(target=_worker, daemon=True)
        self._clf_worker.start()

    def on_reply_piece(self, piece: str, t_reply: float = 0.0) -> None:
        """Feed the streaming reply text (called once per text token)."""
        if self.mode != "autotext" or not self.enable or not piece:
            return
        import time as _time
        now = _time.perf_counter()
        # PersonaPlex text often has NO punctuation, so a pause is the only
        # sentence boundary.  Use WALL-CLOCK time: self.t only advances once per
        # 2 s chunk (after this check), so comparing against it never fires.
        if (self.pending_text.strip() and self._last_piece_wall is not None
                and (now - self._last_piece_wall) > 0.9):
            self._finalize_sentence()
        self.pending_text += piece
        self._last_piece_t = self.t
        self._last_piece_wall = now
        s = self.pending_text.strip()
        if s and s[-1] in ".!?…":
            self._finalize_sentence()
        elif len(self.pending_text) > 120:
            self._finalize_sentence()

    def _finalize_sentence(self) -> None:
        text = self.pending_text.strip()
        self.pending_text = ""
        if len(text) < 3:
            return
        # classify in the background so the render thread never blocks
        self._ensure_clf_worker()
        self._clf_worker_q.put((text, float(self.t)))

    def _classify_text(self, text: str):
        # 1) trained classifier (CPU, background thread)
        self._load_clf()
        if type(self)._CLF is not None:
            try:
                res = type(self)._CLF(text[:480])[0]      # list of {label, score}
                best = max(res, key=lambda r: float(r.get("score", 0.0)))
                lab = str(best.get("label", "")).lower()
                p = float(best.get("score", 0.0))
                cls = self._HF_MAP.get(lab)
                if cls is not None:
                    # guard against noisy short-sentence predictions: if the
                    # winner is weak or neutral is close, stay Neutral
                    neutral_p = 0.0
                    for r in res:
                        if str(r.get("label", "")).lower() == "neutral":
                            neutral_p = float(r.get("score", 0.0))
                    if p < 0.55 or neutral_p > (p - 0.12):
                        cls = 4
                    inten = max(0.22, min(0.75, 0.22 + 0.55 * max(p, neutral_p if cls == 4 else p)))
                    if "!" in text:
                        inten = min(0.85, inten + 0.08)
                    return cls, min(inten, self.inten_max)
            except Exception as e:
                print(f"[emo-autotext] classifier error: {e}", flush=True)
        # 2) lexicon fallback
        low = " " + text.lower() + " "
        for cls in self._LEX_ORDER:
            for kw in self._LEX[cls]:
                if kw in low:
                    inten = 0.45
                    if "!" in text:
                        inten += 0.1
                    letters = [ch for ch in text if ch.isalpha()]
                    if letters and text.upper() == text:
                        inten += 0.1
                    return cls, min(inten, self.inten_max)
        return 4, 0.25

    def _autotext_window(self, t0: float, n: int):
        """Per-frame (one-hot, intensity) for [t0, t0+n/fps).

        A turn is a *state* that stays active until the next turn.  Turns that
        started before this window must still colour the whole window; the
        previous implementation used negative slice indices for those and
        silently fell back to Neutral@0 (v35: lips fully closed).
        """
        if not self.turns:
            return None
        lk = self.lookahead_ms / 1000.0
        fps = self.fps
        if self.neutral_as_happy and self.active:
            cur_c, cur_i = 3, 0.25
        else:
            cur_c, cur_i = 4, 0.0
        emo_f = np.zeros((n, 8), np.float32)
        emo_f[:, cur_c] = 1.0
        inten_f = np.full(n, cur_i, np.float32)

        for (s0, c, it) in self.turns:
            st = max(0.0, s0 - lk)                  # pre-ramp before the turn
            f0 = int(np.floor((st - t0) * fps))
            if f0 >= n:
                break
            f1 = f0 + self.blend
            vec = np.zeros(8, np.float32)
            vec[c] = 1.0
            prev = np.zeros(8, np.float32)
            prev[cur_c] = 1.0
            lo, hi = max(0, f0), min(n, f1)
            if hi > lo:
                a = np.array(
                    [_smoothstep((j - f0 + 1) / (self.blend + 1)) for j in range(lo, hi)],
                    np.float32,
                )
                emo_f[lo:hi] = (1 - a)[:, None] * prev[None, :] + a[:, None] * vec[None, :]
                inten_f[lo:hi] = (1 - a) * cur_i + a * it
            if hi < n:
                # blend finished (or the turn fully predates this window)
                emo_f[hi:] = vec[None, :]
                inten_f[hi:] = it
            cur_c, cur_i = c, it
        return emo_f, inten_f

    def conditioning(self, n_frames: int) -> dict:
        """Build the per-frame conditioning for this chunk and advance time."""
        n = max(1, int(n_frames))
        if not self.enable:
            self.t += n / self.fps
            return {}

        if self.mode == "autotext":
            # finalize a pending sentence when the text stream pauses (>1.2 s):
            # replies often end without punctuation ("I'm doing good")
            import time as _time
            if (self.pending_text.strip() and self._last_piece_wall is not None
                    and (_time.perf_counter() - self._last_piece_wall) > 0.9):
                self._finalize_sentence()
            win = self._autotext_window(self.t, n) if self.turns else None
            if win is None:
                # No emotion turn yet (classifier warming up / first sentence
                # not finalised). v35's Neutral at intensity 0 CLOSES the lips,
                # so a speaking avatar would look mute. Start from the
                # "speaking neutral" (mild Happy) used by the offline renders.
                if self.neutral_as_happy and self.active:
                    _cls0, _it0 = 3, 0.25
                else:
                    _cls0, _it0 = 4, 0.0
                emo_f = np.zeros((n, 8), np.float32)
                emo_f[:, _cls0] = 1.0
                inten_f = np.full(n, _it0, np.float32)
            else:
                emo_f, inten_f = win
            mid = max(0, n // 2)
            self.cls = int(emo_f[mid].argmax())
            self.inten = float(inten_f[mid])
            self._prev_vec = emo_f[-1].copy()
            self._prev_inten = float(inten_f[-1])
        elif self.mode == "script" and self.script:
            if self.active:
                win = self._schedule_window(self.t, n)
            else:
                win = None
            if win is None:
                emo_f = np.zeros((n, 8), np.float32)
                emo_f[:, 4] = 1.0
                inten_f = np.zeros(n, np.float32)
            else:
                emo_f, inten_f = win
            mid = n // 2
            self.cls = int(emo_f[mid].argmax())
            self.inten = float(inten_f[mid])
        elif self.mode == "manual" and self.manual_cls >= 0:
            emo_f = np.zeros((n, 8), np.float32)
            emo_f[:, self.manual_cls] = 1.0
            inten_f = np.full(n, self.manual_inten, np.float32)
            self.cls, self.inten = self.manual_cls, self.manual_inten
        else:
            neutral = self.mode == "neutral" or not self.active
            cls = 4 if neutral else int(self.cls)
            inten = 0.0 if neutral else float(self.inten)
            onehot = np.zeros(8, np.float32)
            onehot[cls] = 1.0
            emo_f = np.repeat(onehot[None, :], n, axis=0)
            inten_f = np.full(n, inten, np.float32)
            k = min(self.ramp, n)
            if k > 0:
                a = np.array(
                    [_smoothstep((i + 1) / (k + 1)) for i in range(k)], np.float32
                )
                emo_f[:k] = (
                    (1.0 - a)[:, None] * self._prev_vec[None, :]
                    + a[:, None] * onehot[None, :]
                )
                inten_f[:k] = (1.0 - a) * self._prev_inten + a * inten
            self.cls, self.inten = int(cls), float(inten)
            self._prev_vec, self._prev_inten = onehot.copy(), inten

        self.t += n / self.fps
        return {
            "emo": emo_f,
            "inten": inten_f,
            "emo_scale": self.emo_scale,
            "static": np.array([1.0], np.float32),
            "static_scale": 1.0,
            "_cls": int(self.cls),
            "_inten": float(self.inten),
        }


class LiveHeliumFMOptions(BaseOptions):
    def initialize(self, parser: argparse.ArgumentParser) -> argparse.ArgumentParser:
        parser = super().initialize(parser)
        parser.set_defaults(wav2vec_sec=0.96)
        parser.add_argument("--host", default="0.0.0.0")
        parser.add_argument("--port", type=int, default=8998)
        parser.add_argument("--html_path", default=str(ROOT / "static" / "index_v3_binary_fullscreen.html"))
        parser.add_argument("--generator_path", required=True)
        parser.add_argument("--lora_generator_path", default="", help="Optional LoRA generator checkpoint to apply on top of --generator_path")
        parser.add_argument("--lora_rank", type=int, default=64)
        parser.add_argument("--lora_alpha", type=float, default=128.0)
        parser.add_argument("--lora_dropout", type=float, default=0.05)
        parser.add_argument("--no_lora_pose_projection", action="store_true")
        parser.add_argument("--no_lora_audio_projection", action="store_true")
        parser.add_argument("--only_lora_pose_projection", action="store_true")
        parser.add_argument("--renderer_path", required=True)
        parser.add_argument("--adapter_path", default="", help="Frontend fp32 Helium->Wav2Vec2 projected-frontend adapter checkpoint (not needed with --audio_direct 1)")
        parser.add_argument(
            "--adapter_type",
            choices=("frontend", "unitalk_last_layer"),
            default="frontend",
            help="Adapter architecture. UniTalk passes only layer 12 to IMTalker.",
        )
        parser.add_argument(
            "--adapter_window_mode",
            choices=("tail", "lookahead"),
            default="tail",
        )
        parser.add_argument(
            "--adapter_future_steps",
            type=int,
            default=6,
            help="12.5 Hz future-context steps retained in lookahead mode.",
        )
        parser.add_argument("--adapter_num_layers", type=int, default=6, help="Transformer layers in the frontend adapter checkpoint")
        parser.add_argument("--adapter_dropout", type=float, default=0.1, help="Dropout value used when constructing the frontend adapter")
        parser.add_argument("--stats_path", default="", help="Unused for frontend adapter mode; accepted for compatibility")
        parser.add_argument(
            "--silence_helium_path",
            default="",
            help="Optional real-silence Helium mean used to initialize the deque.",
        )
        parser.add_argument("--ref_path", required=True)
        parser.add_argument("--audio_path", default="", help="WAV to stream in fixed chunks (simulate-live mode)")
        # Moshi
        parser.add_argument("--moshi_root", default="/workspace/moshi")
        parser.add_argument("--mimi_hf_repo", default="kyutai/moshiko-pytorch-bf16")
        parser.add_argument("--moshi_weight", default="", help="Optional local PersonaPlex/Moshi LM checkpoint")
        parser.add_argument("--mimi_weight", default="", help="Optional local Mimi checkpoint")
        parser.add_argument("--tokenizer", default="", help="Optional local sentencepiece tokenizer")
        parser.add_argument("--quantize_4bit", action="store_true", help="Load PersonaPlex/Moshi LM with bnb 4-bit quantization")
        parser.add_argument("--num_codebooks", type=int, default=8, help="PersonaPlex/Moshi audio codebooks")
        parser.add_argument("--moshi_context", type=int, default=0, help="Optional PersonaPlex/Moshi KV context length")
        parser.add_argument("--voice_prompt", default="", help="PersonaPlex voice prompt filename, e.g. NATM0.pt")
        parser.add_argument("--voice_prompt_dir", default="", help="Optional PersonaPlex voice prompt directory")
        parser.add_argument("--text_prompt", default="", help="Optional PersonaPlex system text prompt")
        parser.add_argument("--moshi_reply_device", default=None, help="Optional separate CUDA device for Moshi reply generation")
        parser.add_argument("--enable_moshi_reply", action="store_true", help="Mic -> Moshi reply audio -> Helium/FM avatar")
        parser.add_argument("--moshi_cfg_coef", type=float, default=1.0)
        parser.add_argument("--direct_reply_hidden", action="store_true", default=True, help="Use Moshi generation hidden directly instead of re-encoding reply audio")
        parser.add_argument("--no_direct_reply_hidden", dest="direct_reply_hidden", action="store_false")
        parser.add_argument(
            "--disable_assistant_output_gate",
            action="store_true",
            help="Disable assistant-output RMS gating and always drive IMTalker from live PersonaPlex hidden states",
        )
        parser.add_argument(
            "--assistant_speech_rms_threshold",
            type=float,
            default=0.006,
            help="RMS threshold on PersonaPlex assistant reply audio before hidden states are allowed to drive avatar motion",
        )
        parser.add_argument(
            "--assistant_speech_hold_chunks",
            type=int,
            default=1,
            help="Keep assistant gate open for this many avatar chunks after reply audio drops below threshold",
        )
        # FM
        parser.add_argument("--audio_chunk_sec", type=float, default=0.96)
        parser.add_argument("--fm_chunk_frames", type=int, default=24, help="Must match wav2vec_sec×fps")
        parser.add_argument(
            "--enable_live_sliding_window",
            action="store_true",
            help="Run IMTalker/FM on a past/current/future feature window and emit only the current slice.",
        )
        parser.add_argument(
            "--slide_past_frames",
            type=int,
            default=10,
            help="Past 25Hz frames used by --enable_live_sliding_window.",
        )
        parser.add_argument(
            "--slide_future_frames",
            type=int,
            default=3,
            help="Future 25Hz frames used by --enable_live_sliding_window; 3 frames = 120ms.",
        )
        parser.add_argument("--skip_fm_audio_encoder", action="store_true", help="Skip FMGenerator's raw-audio Wav2Vec encoder; live PersonaPlex passes precomputed adapter features")
        # emotion conditioning (audio + semantic) for the live IMTalker generator
        parser.add_argument("--emo_enable", type=int, default=1, help="1 = emotion conditioning on")
        parser.add_argument("--emo_mode", default="auto", choices=["auto", "manual", "neutral", "script", "autotext"])
        parser.add_argument("--emo_class", type=int, default=-1, help="manual emotion class 0..7 (-1 = off)")
        parser.add_argument("--emo_inten", type=float, default=0.5, help="manual/script emotion intensity 0..1")
        parser.add_argument("--emo_inten_max", type=float, default=0.55,
                            help="Upper bound for autotext emotion intensity (live)")
        parser.add_argument("--emo_gain", type=float, default=1.0, help="auto-mode intensity multiplier")
        parser.add_argument("--emo_ema", type=float, default=0.35, help="per-chunk emotion EMA (0..1)")
        parser.add_argument("--emo_script", default="",
                            help="scheduled emotions in reply-audio seconds, e.g. "
                                 "'happy@0:0.35,neutral@5:0,sad@6:0.45'. Optional "
                                 "':intensity' overrides --emo_inten.")
        parser.add_argument("--neutral_as_happy", type=int, default=1,
                            help="1 = render Neutral as a mild Happy (0.22-0.30): v35's Neutral "
                                 "conditioning closes the lips completely")
        parser.add_argument("--emo_lookahead_ms", type=float, default=400.0,
                            help="autotext: pre-ramp each detected turn this much earlier "
                                 "(text tags are known at sentence end)")
        parser.add_argument("--head_pin", type=float, default=0.6,
                            help="remove the slow latent component (head drift) in FM.sample; "
                                 "0 disables. Helps between emotion changes and during silence")
        parser.add_argument("--emo_scale", type=float, default=1.0,
                            help="multiplier on the emotion embedding (1.0 = training scale)")
        parser.add_argument("--emo_blend_frames", type=int, default=15,
                            help="smoothstep one-hot blend across an emotion switch "
                                 "(15 = 0.6s at 25fps, same as the offline renders)")
        parser.add_argument("--emo_ramp_frames", type=int, default=12,
                            help="frames over which intensity eases in/out "
                                 "(12 = 0.48s at 25fps, same as the offline renders)")
        parser.add_argument("--e_cfg_scale", type=float, default=3.0, help="emotion classifier-free guidance scale")
        parser.add_argument("--audio_direct", type=int, default=0,
                            help="1 = feed the reply PCM straight into the IMTalker wav2vec encoder "
                                 "(no Helium, no Helium->wav2vec adapter); 0 = legacy Helium path")
        parser.add_argument("--static_pose_cam", action="store_true", default=True, help="pin pose/cam to the measured static values")
        parser.add_argument("--no_static_pose_cam", dest="static_pose_cam", action="store_false")
        parser.add_argument("--reply_hidden_steps_per_chunk", type=int, default=0, help="Raw Moshi 12.5Hz hidden steps per avatar chunk; 0 derives from fm_chunk_frames/fps")
        parser.add_argument("--drain_chunks", type=int, default=1,
                            help="silence chunks emitted at each reply start (video pipeline warm-up; "
                                 "the reply audio is buffered, not dropped)")
        parser.add_argument("--prebuffer_chunks", type=int, default=3, help="Avatar chunks queued before sender starts pacing")
        parser.add_argument("--frame_q_backpressure", type=int, default=160)
        parser.add_argument(
            "--file_chunk_lookahead",
            type=int,
            default=0,
            help="For --audio_path mode, wait this many future chunks before emitting the oldest chunk",
        )
        parser.add_argument("--render_sub_batch", type=int, default=8)
        parser.add_argument("--idle_still", type=int, default=1,
                            help="Render a still, closed-mouth reference pose for silent chunks")
        parser.add_argument("--idle_rms_threshold", type=float, default=0.004,
                            help="Chunk RMS below this is treated as silence")
        parser.add_argument("--idle_crossfade_frames", type=int, default=10)
        parser.add_argument("--motion_crossfade_frames", type=int, default=5,
                            help="frames blended at each chunk boundary to hide the "
                                 "cold-regeneration seam (0 disables)")
        parser.add_argument("--jpeg_quality", type=int, default=86)
        parser.add_argument("--reply_audio_gain", type=float, default=1.0, help="Accepted for launch-script compatibility")
        parser.add_argument(
            "--motion_ref_blend",
            type=float,
            default=0.0,
            help="Blend generated motion latent toward the reference image latent to reduce head tilt.",
        )
        parser.add_argument("--device", default="cuda")
        parser.add_argument("--buffer_ms", type=int, default=80)
        parser.add_argument(
            "--renderer_precision",
            choices=("fp32", "fp16", "bf16"),
            default="fp32",
        )
        parser.add_argument(
            "--output_audio_codec",
            choices=("pcm", "opus"),
            default="pcm",
            help="Assistant audio transport; Opus is one persistent stream per session.",
        )
        parser.add_argument("--dump_motion", action="store_true", help="Dump last session motion/audio to disk")
        parser.add_argument("--dump_dir", default=str(ROOT / "live_try_dumps"))
        # Shared noise
        parser.add_argument("--shared_noise", action="store_true")
        parser.add_argument("--noise_seed", type=int, default=1234)
        parser.add_argument("--noise_max_frames", type=int, default=5000)
        parser.add_argument(
            "--noise_temporal_corr",
            type=float,
            default=0.0,
            help="AR(1)-style temporal correlation applied to FM initial noise.",
        )
        parser.add_argument(
            "--motion_prior_noise_blend",
            type=float,
            default=0.0,
            help="Small blend from previous/generated motion into FM initial noise.",
        )
        # Precision
        parser.add_argument("--fp32", action="store_true")
        parser.add_argument("--tf32", action="store_true")
        parser.add_argument("--compile_renderer", action="store_true")
        # Eye blink motion-map compositing
        parser.add_argument("--enable_eye_blink_composite", action="store_true")
        parser.add_argument("--blink_motion_path", default="", help="Cached blink motion latent .pt file")
        parser.add_argument("--eye_left_x", type=float, default=0.36)
        parser.add_argument("--eye_right_x", type=float, default=0.64)
        parser.add_argument("--eye_center_y", type=float, default=0.405)
        parser.add_argument("--eye_radius_x", type=float, default=0.145)
        parser.add_argument("--eye_radius_y", type=float, default=0.070)
        parser.add_argument("--eye_feather", type=float, default=0.10)
        return parser


# ---------------------------------------------------------------------------
# FastAPI app
# ---------------------------------------------------------------------------

def build_app(args: argparse.Namespace) -> FastAPI:
    app = FastAPI(title="IMTalker Helium MotionField Deque FM liveTry")
    assets_dir = ROOT / "static" / "assets"
    if assets_dir.is_dir():
        app.mount("/assets", StaticFiles(directory=str(assets_dir)), name="assets")
    started_at = time.perf_counter()
    html_path = Path(args.html_path)
    engine: LiveHeliumFMEngine | None = LiveHeliumFMEngine(args)
    moshi_engine: MoshiOnlyEngine | None = None

    def get_engine() -> LiveHeliumFMEngine:
        nonlocal engine
        if engine is None:
            engine = LiveHeliumFMEngine(args)
        return engine

    def get_moshi_engine() -> MoshiOnlyEngine:
        nonlocal moshi_engine
        if moshi_engine is None:
            moshi_engine = MoshiOnlyEngineWithHidden(
                moshi_root=args.moshi_root,
                mimi_hf_repo=args.mimi_hf_repo,
                device=getattr(args, "moshi_reply_device", None) or args.device,
                cfg_coef=float(args.moshi_cfg_coef),
                placeholder_jpeg_b64="",
                moshi_weight=getattr(args, "moshi_weight", ""),
                mimi_weight=getattr(args, "mimi_weight", ""),
                tokenizer=getattr(args, "tokenizer", ""),
                quantize_4bit=bool(getattr(args, "quantize_4bit", False)),
                num_codebooks=int(getattr(args, "num_codebooks", 8)),
                context=(int(args.moshi_context) if int(getattr(args, "moshi_context", 0)) > 0 else None),
                voice_prompt=getattr(args, "voice_prompt", ""),
                voice_prompt_dir=getattr(args, "voice_prompt_dir", ""),
                text_prompt=getattr(args, "text_prompt", ""),
            )
        return moshi_engine

    split_sessions: dict[str, dict] = {}

    async def _get_split_media_epoch(session: dict) -> float:
        while not session["prebuffer_ready"].is_set() and not session.get("closed", False):
            await asyncio.sleep(0.01)
        async with session["media_epoch_lock"]:
            if session.get("media_epoch") is None:
                session["media_epoch"] = time.perf_counter() + 0.08
            return float(session["media_epoch"])

    if bool(getattr(args, "enable_moshi_reply", False)):
        print("[liveTryHeliumFM_ws_binary] eager-loading Moshi/PersonaPlex reply engine", flush=True)
        get_moshi_engine()

    @app.get("/")
    async def index():
        if html_path.is_file():
            return FileResponse(
                html_path,
                headers={
                    "Cache-Control": "no-store, no-cache, must-revalidate, max-age=0",
                    "Pragma": "no-cache",
                    "Expires": "0",
                },
            )
        return HTMLResponse(f"<h1>Missing HTML</h1><p>Expected: {html_path}</p>", status_code=500)

    @app.get("/health")
    async def health():
        return JSONResponse({
            "ok": True,
            "stage": "live_moshi_reply_helium_fm" if args.enable_moshi_reply else "live_helium_fm",
            "uptime_sec": round(time.perf_counter() - started_at, 3),
            "loaded": engine is not None,
        })

    @app.websocket("/ws/video")
    async def video_stream(ws: WebSocket):
        await ws.accept()
        session_id = str(ws.query_params.get("session_id", ""))
        session = split_sessions.get(session_id)
        if session is None:
            await ws.send_json({
                "type": "error",
                "message": "unknown or expired video session",
            })
            await ws.close()
            return

        frame_q: asyncio.Queue = session["frame_q"]
        send_start_wall = await _get_split_media_epoch(session)
        frames_sent = 0
        starvation_events = 0
        starve_start: float | None = None
        print(
            f"[AJ][VIDEO] connected session={session_id[:8]} "
            f"queued={frame_q.qsize()}",
            flush=True,
        )
        try:
            while True:
                try:
                    packet = frame_q.get_nowait()
                except asyncio.QueueEmpty:
                    if starve_start is None:
                        starve_start = time.perf_counter()
                        starvation_events += 1
                    await asyncio.sleep(0.004)
                    if session.get("closed", False) and frame_q.empty():
                        break
                    continue

                if packet is None:
                    break

                if starve_start is not None:
                    gap_ms = 1000.0 * (time.perf_counter() - starve_start)
                    if gap_ms > 100:
                        print(
                            f"[AJ][VIDEO] STARVED {gap_ms:.0f}ms "
                            f"(event #{starvation_events}) "
                            f"frame_q={frame_q.qsize()} sent={frames_sent}",
                            flush=True,
                        )
                    starve_start = None

                idx = int(packet["frame_number"])
                target_t = send_start_wall + idx / float(args.fps)
                sleep_s = target_t - time.perf_counter()
                if sleep_s > 0:
                    await asyncio.sleep(sleep_s)

                if packet.get("ws_kind") == "bytes":
                    await ws.send_bytes(packet["data"])
                else:
                    await ws.send_json(packet["msg"])
                frames_sent += 1

                late_s = time.perf_counter() - target_t
                if late_s > 0.5:
                    print(
                        f"[AJ][VIDEO] frame {idx} is {late_s*1000:.0f}ms late",
                        flush=True,
                    )
                if frames_sent % 50 == 0:
                    elapsed = time.perf_counter() - send_start_wall
                    print(
                        f"[AJ][VIDEO] sent={frames_sent} frame={idx} "
                        f"frame_q={frame_q.qsize()} elapsed={elapsed:.1f}s "
                        f"starve_events={starvation_events}",
                        flush=True,
                    )
        except (WebSocketDisconnect, RuntimeError, Exception) as exc:
            print(f"[AJ][VIDEO] closed session={session_id[:8]} {exc!r}", flush=True)
        finally:
            print(f"[AJ][VIDEO] done session={session_id[:8]} sent={frames_sent}", flush=True)

    @app.websocket("/ws/conversation")
    async def conversation(ws: WebSocket):
        await ws.accept()
        fm_engine = get_engine()
        fm_engine.reset_session()
        reply_engine = get_moshi_engine() if args.enable_moshi_reply else None
        if reply_engine is not None:
            reply_engine.reset_session()
        browser_input_sr = 48000
        opus_reader = sphn.OpusStreamReader(TARGET_SR)
        audio_packets_seen = 0

        # -- Queues for the reply pipeline --
        # mic_q: (raw_bytes, input_sr) or None to stop
        # frame_q: per-frame packet dicts or None to stop
        mic_q: queue.Queue[tuple[bytes, int] | None] | None = None
        frame_q: asyncio.Queue[dict | None] | None = None
        audio_q: asyncio.Queue[dict | None] | None = None
        gpu_thread: threading.Thread | None = None
        sender_task: asyncio.Task | None = None
        audio_sender_task: asyncio.Task | None = None
        event_loop: asyncio.AbstractEventLoop | None = None
        ws_send_lock = asyncio.Lock()
        media_epoch_lock = asyncio.Lock()
        media_epoch: float | None = None
        session_id = uuid.uuid4().hex
        prebuffer_ready = threading.Event()
        session_started = threading.Event()
        last_mic_level_log_wall = 0.0

        stream_task: asyncio.Task | None = None

        if reply_engine is not None:
            mic_q = queue.Queue(maxsize=512)
            frame_q = asyncio.Queue(maxsize=512)
            audio_q = asyncio.Queue(maxsize=512)
            event_loop = asyncio.get_running_loop()
            split_sessions[session_id] = {
                "frame_q": frame_q,
                "prebuffer_ready": prebuffer_ready,
                "media_epoch_lock": media_epoch_lock,
                "media_epoch": media_epoch,
                "closed": False,
                "created_at": time.perf_counter(),
            }

        await ws.send_json({
            "type": "server_ready",
            "variant": (
                "AJ-NETWORK-ISO-CACHED-FP32"
                if os.environ.get("IMTALKER_CACHED_ENGINE", "0").strip().lower()
                in {"1", "true", "yes", "on"}
                else "AJ-NETWORK-ISO-FP32"
            ),
            "session_id": session_id,
            "video_ws_path": f"/ws/video?session_id={session_id}",
            "sample_rate": TARGET_SR,
            "output_audio_codec": str(args.output_audio_codec),
            "model_type": "moshi_reply+helium_fm+renderer" if args.enable_moshi_reply else "helium_fm+renderer",
            "tokens_per_chunk": int(args.fm_chunk_frames),
            "has_audio_file": fm_engine.audio_pcm is not None,
            "buffer_ms": int(args.buffer_ms),
            "av_transport": "binary",
            "target_fps": round(float(args.fps), 2),
        })
        print("[liveTryHeliumFM] websocket connected; sent server_ready", flush=True)

        if reply_engine is not None:
            def _gpu_producer_thread() -> None:
                """GPU thread: Moshi -> Helium -> FM -> render -> JPEG -> frame_q.

                Runs Moshi at maximum GPU speed. When real mic audio is not
                yet available, pads Moshi input with silence so it can keep
                generating reply audio without waiting for real-time mic
                arrival. This dramatically cuts first-reply latency.
                """
                assert (
                    mic_q is not None
                    and frame_q is not None
                    and audio_q is not None
                    and event_loop is not None
                )
                pending_reply_steps: list[dict] = []
                pending_reply_hidden: list[torch.Tensor] = []
                pending_reply_audio: list[np.ndarray] = []
                reply_step_history: list[dict] = []
                reply_audio_history: list[np.ndarray] = []
                reply_avatar_chunk_idx = 0
                chunk_produce_count = 0
                stopped = False
                # Pause producer when queue is this deep (qsize() is heuristic across threads).
                FRAME_Q_BACKPRESS = max(1, int(getattr(args, "frame_q_backpressure", 96)))
                FRAME_Q_PUT_TIMEOUT_S = 120.0
                prebuffer_chunks = max(0, int(getattr(args, "prebuffer_chunks", PREBUFFER_CHUNKS)))
                hidden_steps_per_chunk = int(getattr(args, "reply_hidden_steps_per_chunk", 0))
                if hidden_steps_per_chunk <= 0:
                    hidden_steps_per_chunk = int(round(float(args.fm_chunk_frames) * 12.5 / float(args.fps)))
                hidden_steps_per_chunk = max(1, hidden_steps_per_chunk)
                max_moshi_steps_per_loop = 1
                was_silent = True
                assistant_gate_hold = 0
                last_motion_frame: torch.Tensor | None = None

                def _enqueue_frame(pkt: dict) -> None:
                    """Block until frame_q accepts pkt (real backpressure). Must run from GPU thread."""
                    fut = asyncio.run_coroutine_threadsafe(frame_q.put(pkt), event_loop)
                    try:
                        fut.result(timeout=FRAME_Q_PUT_TIMEOUT_S)
                    except TimeoutError:
                        print(
                            f"[GPU] WARNING frame_q.put timeout frame={pkt.get('frame_number')}",
                            flush=True,
                        )
                    except Exception as e:
                        print(f"[GPU] WARNING frame_q.put failed: {e!r}", flush=True)

                def _enqueue_audio(pkt: dict) -> None:
                    # NEVER block the GPU thread on the legacy audio queue: the
                    # binary AV frames already carry the audio, and the aj_nodrop
                    # client never opens a separate audio websocket. Without a
                    # consumer this queue used to fill at 512 packets and stall
                    # production for seconds -> the avatar froze after ~20 s.
                    def _put():
                        try:
                            audio_q.put_nowait(pkt)
                        except asyncio.QueueFull:
                            try:
                                audio_q.get_nowait()   # drop the oldest
                            except asyncio.QueueEmpty:
                                pass
                            try:
                                audio_q.put_nowait(pkt)
                            except asyncio.QueueFull:
                                pass
                    try:
                        event_loop.call_soon_threadsafe(_put)
                    except Exception as e:
                        print(f"[GPU] WARNING audio_q.put failed: {e!r}", flush=True)
                    except Exception as e:
                        print(f"[GPU] WARNING audio_q.put failed: {e!r}", flush=True)

                if prebuffer_chunks <= 0 and not prebuffer_ready.is_set():
                    prebuffer_ready.set()
                    print("[GPU] prebuffer=0, sender starts immediately", flush=True)

                while not stopped:
                    # --- Phase 1: drain all available mic audio (non-blocking) ---
                    while True:
                        try:
                            item = mic_q.get_nowait()
                        except queue.Empty:
                            break
                        if item is None:
                            stopped = True
                            break
                        raw_bytes, input_sr = item
                        if raw_bytes:
                            reply_engine.append_browser_pcm(
                                np.frombuffer(raw_bytes, dtype=np.int16), input_sr
                            )

                    if stopped:
                        break

                    # --- Phase 2: yield if outbound queue is deep (sender is the bottleneck) ---
                    while frame_q.qsize() >= FRAME_Q_BACKPRESS:
                        time.sleep(0.004)

                    if not session_started.is_set():
                        time.sleep(0.003)
                        continue

                    if reply_engine.input_buffer.shape[0] < MIMI_FRAME_SIZE:
                        time.sleep(0.003)
                        continue

                    # --- Phase 3: run all Moshi steps and feed FM ---
                    t_recv = time.perf_counter()
                    for ev in reply_engine.process_ready_steps_limited(max_moshi_steps_per_loop):
                        pending_reply_steps.append(ev)
                        fm_engine._session_reply_events.append({
                            "step": int(ev.get("step", -1)),
                            "token": int(ev.get("token", -1)),
                            "piece": str(ev.get("piece", "")),
                            "audio_text": str(ev.get("audio_text", "")),
                            "reply_rms": float(ev.get("reply_rms", 0.0)),
                            "reply_peak": float(ev.get("reply_peak", 0.0)),
                            "input_rms": float(ev.get("input_rms", 0.0)),
                            "hidden": bool(isinstance(ev.get("helium_hidden"), torch.Tensor)),
                            "total_ms": float(ev.get("total_ms", 0.0)),
                        })
                        if str(ev.get("piece", "")):
                            try:
                                fm_engine.emo.on_reply_piece(str(ev.get("piece", "")))
                            except Exception as _e:
                                print(f"[emo-autotext] ingest error {_e}", flush=True)
                        reply_pcm = (
                            np.frombuffer(base64.b64decode(ev["reply_i16_b64"]), dtype=np.int16)
                            .astype(np.float32) / 32768.0
                        )
                        if bool(getattr(args, "direct_reply_hidden", False)):
                            hidden = ev.get("helium_hidden")
                            if not isinstance(hidden, torch.Tensor):
                                continue
                            pending_reply_hidden.append(hidden.squeeze(0).contiguous())
                            pending_reply_audio.append(reply_pcm)
                            if len(pending_reply_hidden) < hidden_steps_per_chunk:
                                continue

                            used_hidden = pending_reply_hidden[:hidden_steps_per_chunk]
                            used_audio = pending_reply_audio[:hidden_steps_per_chunk]
                            used_steps = pending_reply_steps[:hidden_steps_per_chunk]
                            pending_reply_hidden = pending_reply_hidden[hidden_steps_per_chunk:]
                            pending_reply_audio = pending_reply_audio[hidden_steps_per_chunk:]
                            pending_reply_steps = pending_reply_steps[hidden_steps_per_chunk:]

                            if str(getattr(args, "adapter_window_mode", "tail")) == "lookahead":
                                history_size = int(fm_engine.helium_deque_size)
                                future_steps = int(
                                    getattr(args, "adapter_future_steps", 6)
                                )
                                reply_audio_history.extend(used_audio)
                                reply_step_history.extend(used_steps)
                                reply_audio_history = reply_audio_history[-history_size:]
                                reply_step_history = reply_step_history[-history_size:]

                                zero_audio = np.zeros_like(used_audio[0])
                                audio_window = (
                                    [zero_audio]
                                    * (history_size - len(reply_audio_history))
                                    + reply_audio_history
                                )
                                silent_step = {
                                    "token": 0,
                                    "reply_rms": 0.0,
                                    "total_ms": 0.0,
                                    "audio_text": "",
                                }
                                step_window = (
                                    [silent_step]
                                    * (history_size - len(reply_step_history))
                                    + reply_step_history
                                )
                                output_end = history_size - future_steps
                                output_start = (
                                    output_end - hidden_steps_per_chunk
                                )
                                used_audio = audio_window[output_start:output_end]
                                used_steps = step_window[output_start:output_end]

                            pcm_chunk = np.concatenate(used_audio, axis=0).astype(np.float32, copy=False)
                            reply_rms = float(np.sqrt(np.mean(np.square(pcm_chunk)))) if pcm_chunk.size else 0.0
                            step_rms = max(
                                (
                                    float(s.get("reply_rms", 0.0) or 0.0)
                                    for s in used_steps
                                ),
                                default=0.0,
                            )
                            speech_threshold = float(
                                getattr(args, "assistant_speech_rms_threshold", 0.006)
                            )
                            assistant_active_now = max(reply_rms, step_rms) > speech_threshold
                            if assistant_active_now:
                                assistant_gate_hold = max(
                                    0,
                                    int(getattr(args, "assistant_speech_hold_chunks", 1)),
                                )
                            elif assistant_gate_hold > 0:
                                assistant_gate_hold -= 1
                            assistant_active = assistant_active_now or assistant_gate_hold > 0
                            if bool(getattr(args, "disable_assistant_output_gate", False)):
                                assistant_active = True

                            previous_active = not was_silent
                            transition = assistant_active != previous_active
                            if transition:
                                print(
                                    f"[GPU] Avatar gate transition "
                                    f"{'speech' if assistant_active else 'idle'} "
                                    f"reply_rms={reply_rms:.5f} step_rms={step_rms:.5f}",
                                    flush=True,
                                )
                            was_silent = not assistant_active

                            if assistant_active:
                                helium_chunk = torch.cat(used_hidden, dim=0)
                            else:
                                if fm_engine.silence_helium_seed is not None:
                                    helium_chunk = fm_engine.silence_helium_seed.expand(
                                        hidden_steps_per_chunk, -1
                                    ).contiguous()
                                else:
                                    hidden_dim = int(used_hidden[0].shape[-1])
                                    helium_chunk = torch.zeros(
                                        hidden_steps_per_chunk,
                                        hidden_dim,
                                        device=fm_engine.device,
                                        dtype=torch.float32,
                                    )
                            # ---- DRAIN: give the video pipeline a head start ----
                            # While the first `drain_chunks` of a reply are emitted
                            # as silence, the A/V pipeline warms up; the reply audio
                            # is buffered and played afterwards, so the mouth is
                            # already synced when speech starts (fixes the perceived
                            # "avatar lags behind the audio" at reply start).
                            _drain_n = int(getattr(args, "drain_chunks", 1))
                            _prev_active = bool(getattr(fm_engine, "_was_active", False))
                            if assistant_active and not _prev_active:
                                fm_engine._drain_left = _drain_n
                                fm_engine._drain_buf = []
                                if _drain_n > 0:
                                    print(f"[drain] reply start: {_drain_n} silence chunk(s) "
                                          f"({_drain_n * args.audio_chunk_sec:.1f}s lead-in)",
                                          flush=True)
                            fm_engine._was_active = assistant_active
                            if not hasattr(fm_engine, "_drain_buf") or fm_engine._drain_buf is None:
                                fm_engine._drain_buf = []
                            if not hasattr(fm_engine, "_drain_left"):
                                fm_engine._drain_left = 0
                            if used_audio:
                                fm_engine._drain_buf.append(
                                    np.concatenate(used_audio).astype(np.float32))
                            if int(fm_engine._drain_left) > 0 and fm_engine._drain_buf:
                                fm_engine._drain_left = int(fm_engine._drain_left) - 1
                                _emit_audio = np.zeros(
                                    int(round(args.audio_chunk_sec * TARGET_SR)), np.float32)
                            elif fm_engine._drain_buf:
                                _emit_audio = fm_engine._drain_buf.pop(0)
                            else:
                                _emit_audio = (np.concatenate(used_audio).astype(np.float32)
                                               if used_audio else np.zeros(0, np.float32))
                            if _emit_audio.shape[0] == 0:
                                _emit_audio = np.zeros(
                                    int(round(args.audio_chunk_sec * TARGET_SR)), np.float32)
                            target_frames = max(1, int(round(len(_emit_audio) * float(args.fps) / TARGET_SR)))
                            # CRITICAL: the frames (and the recorded session audio) must
                            # carry exactly the audio the motion was conditioned on.
                            # The drain delays the reply by `drain_chunks`; using the
                            # undelayed `pcm_chunk` here put the audio ~2 s AHEAD of
                            # the lip motion -> the avatar "lagged behind the audio".
                            pcm_chunk = _emit_audio
                            _emit_rms = (float(np.sqrt(np.mean(np.square(_emit_audio))))
                                         if _emit_audio.size else 0.0)
                            _emo_active = bool(
                                assistant_active
                                and _emit_rms > float(getattr(args, "assistant_speech_rms_threshold", 0.006))
                            )
                            fm_engine.emo.update(used_audio, used_steps, _emo_active)
                            fm_engine.pending_audio = (
                                _emit_audio if _emit_audio.shape[0] else None
                            )
                            motion, fm_info = fm_engine._sample_motion_from_helium(helium_chunk, target_frames)
                            # Every chunk is generated independently (cold), so always
                            # blend the first frames onto the previous chunk's last frame.
                            # Without this the boundary shows a visible head/jaw jump.
                            if last_motion_frame is not None:
                                blend_frames = min(
                                    int(getattr(args, "motion_crossfade_frames", 4)),
                                    int(motion.shape[0]),
                                )
                                if blend_frames > 0:
                                    alpha = torch.linspace(
                                        0.0,
                                        1.0,
                                        blend_frames + 2,
                                        device=motion.device,
                                        dtype=motion.dtype,
                                    )[1:-1].unsqueeze(-1)
                                    previous = last_motion_frame.to(
                                        device=motion.device,
                                        dtype=motion.dtype,
                                    ).unsqueeze(0)
                                    motion[:blend_frames] = (
                                        previous * (1.0 - alpha)
                                        + motion[:blend_frames] * alpha
                                    )
                            last_motion_frame = motion[-1].detach().clone()
                            fm_info["assistant_active"] = bool(assistant_active)
                            fm_info["assistant_reply_rms"] = float(reply_rms)
                            fm_info["assistant_step_rms"] = float(step_rms)
                            used_codes = [
                                s["reply_codes"].to(dtype=torch.int16).contiguous()
                                for s in used_steps
                                if isinstance(s.get("reply_codes"), torch.Tensor)
                            ]
                            if used_codes:
                                fm_engine._session_live_token_parts.extend(used_codes)
                            fm_engine._record_session_chunk(pcm_chunk, motion, fm_info)
                        else:
                            result = fm_engine.feed_pcm_f32(reply_pcm)
                            if result is None:
                                continue

                            motion, fm_info, pcm_chunk = result
                            steps_per_avatar_chunk = max(1, int(round(len(pcm_chunk) / MIMI_FRAME_SIZE)))
                            used_steps = pending_reply_steps[:steps_per_avatar_chunk]
                            pending_reply_steps = pending_reply_steps[steps_per_avatar_chunk:]
                            used_codes = [
                                s["reply_codes"].to(dtype=torch.int16).contiguous()
                                for s in used_steps
                                if isinstance(s.get("reply_codes"), torch.Tensor)
                            ]
                            if used_codes:
                                fm_engine._session_live_token_parts.extend(used_codes)

                        reply_avatar_chunk_idx += 1
                        chunk_produce_count += 1

                        moshi_total_ms = sum(float(s.get("total_ms", 0.0)) for s in used_steps)

                        avatar_chunk_id = len(fm_engine._session_chunk_rows)
                        try:
                            if os.environ.get("DUMP_CHUNKS") == "1":
                                os.makedirs("/workspace/dbg_chunks", exist_ok=True)
                                np.save(f"/workspace/dbg_chunks/ck{reply_avatar_chunk_idx-1:02d}.npy",
                                        np.asarray(pcm_chunk, dtype=np.float32))
                        except Exception as _e:
                            print(f"[DUMP] err {_e}", flush=True)
                        if os.environ.get("DUMP_CHUNKS") == "1":
                            try:
                                _fr = float(np.sqrt(np.mean(np.square(pcm_chunk)))) if pcm_chunk.size else 0.0
                                _pa = getattr(fm_engine, "pending_audio", None)
                                _pr = float(np.sqrt(np.mean(np.square(_pa)))) if _pa is not None and _pa.size else -1.0
                                print(f"[FRAME-AUDIO] chunk#{reply_avatar_chunk_idx-1} frame_rms={_fr:.4f} "
                                      f"model_rms={_pr:.4f} n={pcm_chunk.size}", flush=True)
                            except Exception as _e:
                                print(f"[FRAME-AUDIO] err {_e}", flush=True)
                        frame_audio = split_audio_into_frame_slices(pcm_chunk, args.fps)
                        n_frames = int(motion.shape[0])
                        emitted = int(fm_info["abs_start"])
                        output_step = used_steps[-1] if used_steps else ev
                        text_payload = (
                            output_step.get("audio_text")
                            or output_step.get("sampled_text")
                            or ""
                        )
                        total_gen_ms = (
                            moshi_total_ms + float(fm_info["helium_ms"]) + float(fm_info["fm_ms"])
                        )

                        for frame_offset, audio_slice in enumerate(frame_audio[:n_frames]):
                            _enqueue_audio({
                                "frame_number": emitted + frame_offset,
                                "audio_pcm": np.asarray(
                                    audio_slice,
                                    dtype=np.float32,
                                ).copy(),
                            })

                        t_chunk_start = time.perf_counter()

                        for sb_start in range(0, n_frames, fm_engine.render_sub_batch):
                            sb_end = min(sb_start + fm_engine.render_sub_batch, n_frames)
                            sub_motion = motion[sb_start:sb_end]
                            sub_audio = frame_audio[sb_start:sb_end]

                            packets = fm_engine.render_and_encode_subbatch(
                                sub_motion,
                                sub_audio,
                                abs_start=emitted + sb_start,
                                text_payload=text_payload,
                                avatar_chunk_id=avatar_chunk_id,
                                total_gen_ms=total_gen_ms,
                            )

                            for pkt in packets:
                                _enqueue_frame(pkt)

                        chunk_wall_ms = _ms(t_chunk_start)
                        produce_latency_ms = _ms(t_recv)
                        q_depth = frame_q.qsize() if frame_q is not None else -1

                        print(
                            f"[GPU][chunk#{reply_avatar_chunk_idx}] "
                            f"moshi={moshi_total_ms:.0f}ms "
                            f"helium={float(fm_info['helium_ms']):.0f}ms "
                            f"fm={float(fm_info['fm_ms']):.0f}ms "
                            f"render+jpeg={chunk_wall_ms:.0f}ms "
                            f"frames={n_frames} "
                            f"emo={fm_info.get('emo_cls', '-')}/{fm_info.get('emo_inten', '-')} "
                            f"produce_latency={produce_latency_ms:.0f}ms "
                            f"frame_q={q_depth} "
                            f"abs={emitted}",
                            flush=True,
                        )

                        if (
                            chunk_produce_count == prebuffer_chunks
                            and not prebuffer_ready.is_set()
                        ):
                            prebuffer_ready.set()
                            print(
                                f"[GPU] prebuffer ready: {q_depth} frames queued "
                                f"after {chunk_produce_count} chunks",
                                flush=True,
                            )

                fut_done = asyncio.run_coroutine_threadsafe(frame_q.put(None), event_loop)
                try:
                    fut_done.result(timeout=30.0)
                except Exception as e:
                    print(f"[GPU] WARNING sentinel put: {e!r}", flush=True)
                fut_audio_done = asyncio.run_coroutine_threadsafe(
                    audio_q.put(None),
                    event_loop,
                )
                try:
                    fut_audio_done.result(timeout=30.0)
                except Exception as e:
                    print(f"[GPU] WARNING audio sentinel put: {e!r}", flush=True)

            async def _get_media_epoch() -> float:
                nonlocal media_epoch
                while not prebuffer_ready.is_set():
                    await asyncio.sleep(0.01)
                async with media_epoch_lock:
                    session = split_sessions.get(session_id)
                    if media_epoch is None and session is not None and session.get("media_epoch") is not None:
                        media_epoch = float(session["media_epoch"])
                    if media_epoch is None:
                        media_epoch = time.perf_counter() + 0.08
                        if session is not None:
                            session["media_epoch"] = media_epoch
                    return media_epoch

            async def _reply_sender() -> None:
                """Wait for prebuffer, then drain frame_q at 25fps."""
                assert frame_q is not None
                send_start_wall = await _get_media_epoch()
                frames_sent = 0
                starvation_events = 0
                starve_start: float | None = None
                ws_closed = False

                q_depth_at_start = frame_q.qsize()
                print(
                    f"[SENDER] prebuffer filled, starting pacing with "
                    f"{q_depth_at_start} frames queued",
                    flush=True,
                )

                while True:
                    if ws_closed:
                        break

                    try:
                        packet = frame_q.get_nowait()
                    except asyncio.QueueEmpty:
                        if send_start_wall is not None and starve_start is None:
                            starve_start = time.perf_counter()
                            starvation_events += 1
                        await asyncio.sleep(0.004)
                        continue

                    if packet is None:
                        break

                    if starve_start is not None:
                        gap_ms = 1000.0 * (time.perf_counter() - starve_start)
                        if gap_ms > 100:
                            print(
                                f"[SENDER] STARVED {gap_ms:.0f}ms "
                                f"(event #{starvation_events}) "
                                f"frame_q={frame_q.qsize()} sent={frames_sent}",
                                flush=True,
                            )
                        starve_start = None

                    idx = int(packet["frame_number"])

                    target_t = send_start_wall + idx / float(args.fps)
                    sleep_s = target_t - time.perf_counter()
                    if sleep_s > 0:
                        await asyncio.sleep(sleep_s)

                    try:
                        async with ws_send_lock:
                            if packet.get("ws_kind") == "bytes":
                                await ws.send_bytes(packet["data"])
                            else:
                                await ws.send_json(packet["msg"])
                    except (WebSocketDisconnect, RuntimeError, Exception):
                        ws_closed = True
                        break
                    frames_sent += 1

                    late_s = time.perf_counter() - target_t
                    if late_s > 0.5:
                        print(
                            f"[SENDER] video frame {idx} is {late_s*1000:.0f}ms late",
                            flush=True,
                        )

                    if frames_sent % 50 == 0:
                        q_depth = frame_q.qsize()
                        elapsed = time.perf_counter() - send_start_wall
                        print(
                            f"[SENDER] sent={frames_sent} frame={idx} "
                            f"frame_q={q_depth} "
                            f"elapsed={elapsed:.1f}s "
                            f"starve_events={starvation_events}",
                            flush=True,
                        )

            async def _audio_sender() -> None:
                assert audio_q is not None
                output_audio_codec = str(args.output_audio_codec).lower()
                if output_audio_codec != "opus":
                    raise RuntimeError("AJ requires --output_audio_codec opus")
                opus_writer = sphn.OpusStreamWriter(TARGET_SR)
                send_start_wall = await _get_media_epoch()
                packets_sent = 0
                bytes_sent = 0
                # AG: prevent catch-up bursts after GPU/render stalls.
                # Even if target_t is already late, keep websocket audio writes
                # spaced near the media frame cadence instead of flushing a
                # backlog back-to-back into the browser decoder/worklet.
                min_send_interval_s = max(0.001, 0.92 / float(args.fps))
                next_send_wall = send_start_wall

                while True:
                    packet = await audio_q.get()
                    if packet is None:
                        break
                    idx = int(packet["frame_number"])
                    target_t = send_start_wall + idx / float(args.fps)
                    target_t = max(target_t, next_send_wall)
                    sleep_s = target_t - time.perf_counter()
                    if sleep_s > 0:
                        await asyncio.sleep(sleep_s)
                    next_send_wall = max(target_t, time.perf_counter()) + min_send_interval_s

                    audio_pcm = np.asarray(
                        packet["audio_pcm"],
                        dtype=np.float32,
                    ).reshape(-1)
                    opus_payload = opus_writer.append_pcm(audio_pcm)
                    if opus_payload is None and hasattr(opus_writer, "read_bytes"):
                        opus_payload = opus_writer.read_bytes()
                    if opus_payload:
                        try:
                            async with ws_send_lock:
                                await ws.send_bytes(b"\x01" + opus_payload)
                        except (WebSocketDisconnect, RuntimeError, Exception):
                            break
                        packets_sent += 1
                        bytes_sent += len(opus_payload)

                    if packets_sent and packets_sent % 50 == 0:
                        print(
                            f"[AUDIO] sent={packets_sent} frame={idx} "
                            f"audio_q={audio_q.qsize()} bytes={bytes_sent}",
                            flush=True,
                        )

            gpu_thread = threading.Thread(target=_gpu_producer_thread, daemon=True, name="gpu-producer")
            gpu_thread.start()
            print(
                f"[AJ][AUDIO] session={session_id[:8]} "
                f"video_path=/ws/video?session_id={session_id}",
                flush=True,
            )
            audio_sender_task = asyncio.create_task(_audio_sender())

        try:
            while True:
                msg = await ws.receive()
                if msg["type"] == "websocket.disconnect":
                    break
                data = msg.get("bytes")
                if data is not None:
                    if reply_engine is None:
                        continue
                    assert mic_q is not None
                    packet_sr = int(browser_input_sr)
                    raw = bytes(data)
                    is_ogg = raw[:4] == b"OggS" or raw[1:5] == b"OggS"
                    if is_ogg:
                        payload = raw[1:] if raw[1:5] == b"OggS" else raw
                        try:
                            decoded_pcm = opus_reader.append_bytes(payload)
                        except Exception as _e:
                            print(f"[MIC] opus decode error (skipped): {_e}", flush=True)
                            continue
                        if decoded_pcm.shape[-1] == 0:
                            continue
                        decoded_pcm = np.asarray(decoded_pcm, dtype=np.float32).reshape(-1)
                        decoded_i16 = (
                            np.clip(decoded_pcm, -1.0, 1.0) * 32767.0
                        ).astype(np.int16)
                        data = decoded_i16.tobytes()
                        packet_sr = TARGET_SR
                    audio_packets_seen += 1
                    pcm_i16 = np.frombuffer(data, dtype=np.int16)
                    mic_rms = float(np.sqrt(np.mean((pcm_i16.astype(np.float32) / 32768.0) ** 2))) if pcm_i16.size else 0.0
                    mic_peak = float(np.max(np.abs(pcm_i16.astype(np.float32) / 32768.0))) if pcm_i16.size else 0.0
                    now_wall = time.perf_counter()
                    if audio_packets_seen <= 3 or now_wall - last_mic_level_log_wall >= 1.0:
                        voice = "VOICE" if mic_rms >= 0.02 else "quiet"
                        print(
                            f"[MIC] packet={audio_packets_seen} rms={mic_rms:.5f} "
                            f"peak={mic_peak:.3f} sr={packet_sr} {voice}",
                            flush=True,
                        )
                        last_mic_level_log_wall = now_wall
                    if audio_packets_seen <= 3:
                        print(
                            f"[liveTryHeliumFM] rx binary mic packet#{audio_packets_seen} "
                            f"bytes={len(data)} sr={packet_sr}",
                            flush=True,
                        )
                    if not session_started.is_set():
                        session_started.set()
                        print(
                            "[liveTryHeliumFM] auto-started session from first binary mic packet",
                            flush=True,
                        )
                    try:
                        mic_q.put_nowait((bytes(data), int(packet_sr)))
                    except queue.Full:
                        pass
                    continue
                text = msg.get("text")
                if text is None:
                    continue
                try:
                    payload = json.loads(text)
                except json.JSONDecodeError:
                    continue
                msg_type = str(payload.get("type", "")).lower()
                print(f"[liveTryHeliumFM] rx text type={msg_type or '<empty>'}", flush=True)

                if msg_type == "start":
                    browser_input_sr = int(payload.get("sample_rate", payload.get("sampleRate", browser_input_sr)))
                    if fm_engine.audio_pcm is not None and (stream_task is None or stream_task.done()):
                        fm_engine.reset_session()
                        stream_task = asyncio.create_task(stream_from_file(ws, fm_engine))
                        def _on_stream_done(t, _ws=ws):
                            if t.cancelled():
                                print("[stream_from_file] cancelled", flush=True)
                                return
                            exc = t.exception()
                            if exc is not None:
                                import traceback as _tb
                                print("[stream_from_file] CRASHED:", repr(exc), flush=True)
                                _tb.print_exception(type(exc), exc, exc.__traceback__)
                        stream_task.add_done_callback(_on_stream_done)
                    if reply_engine is not None:
                        reply_engine.reset_session()
                        session_started.set()
                    # tell the UI the session is armed: the try21 client only
                    # enables its microphone on this message (without it the
                    # user is never heard and no frames are ever produced)
                    try:
                        await ws.send_json({
                            "type": "session_started",
                            "reset_ms": 0.0,
                        })
                        print("[liveTryHeliumFM] sent session_started to client", flush=True)
                    except Exception as _e:
                        print(f"[liveTryHeliumFM] session_started send failed: {_e}", flush=True)
                    print(
                        "[liveTryHeliumFM] start → "
                        + ("streaming from file" if fm_engine.audio_pcm is not None else "live Moshi reply mode"),
                        flush=True,
                    )

                elif msg_type == "chunk_audio":
                    pcm_b64 = payload.get("pcm_s16le_b64", "")
                    if not pcm_b64:
                        continue
                    pcm_bytes = base64.b64decode(pcm_b64)
                    result = fm_engine.feed_pcm(pcm_bytes)
                    if result is not None:
                        motion, fm_info, _pcm_chunk = result
                        avatar_chunk_id = len(fm_engine._session_chunk_rows)
                        n_frames = int(motion.shape[0])
                        emitted = fm_info["abs_start"]
                        for sb_start in range(0, n_frames, fm_engine.render_sub_batch):
                            sub = motion[sb_start:sb_start + fm_engine.render_sub_batch].to(
                                fm_engine.device, dtype=fm_engine.dtype
                            )
                            frames_np, _ = fm_engine._render_motion(sub)
                            for j, frame_rgb in enumerate(frames_np):
                                idx = emitted + sb_start + j
                                await ws.send_json({
                                    "type": "chunk_frame",
                                    "chunk_id": idx + 1,
                                    "frame_idx": 0,
                                    "jpeg_b64": encode_jpeg_b64(frame_rgb, fm_engine.jpeg_quality),
                                    "moshi_text": (
                                        f"live Helium+FM "
                                        f"helium={fm_info['helium_ms']:.0f}ms "
                                        f"fm={fm_info['fm_ms']:.0f}ms"
                                    ),
                                    "server_fps": round(float(args.fps), 1),
                                    "chunks_done": avatar_chunk_id,
                                })

                elif msg_type == "stop":
                    print("[liveTryHeliumFM] stop requested", flush=True)
                    break

        except (WebSocketDisconnect, RuntimeError):
            pass
        finally:
            session = split_sessions.get(session_id)
            if session is not None:
                session["closed"] = True
            if stream_task is not None:
                stream_task.cancel()
                with contextlib.suppress(asyncio.CancelledError):
                    await stream_task
            if mic_q is not None:
                with contextlib.suppress(queue.Full):
                    mic_q.put_nowait(None)
            if gpu_thread is not None:
                await asyncio.to_thread(gpu_thread.join, 30.0)
            if sender_task is not None:
                sender_task.cancel()
                with contextlib.suppress(asyncio.CancelledError, WebSocketDisconnect, RuntimeError):
                    await sender_task
            if audio_sender_task is not None:
                audio_sender_task.cancel()
                with contextlib.suppress(asyncio.CancelledError, WebSocketDisconnect, RuntimeError):
                    await audio_sender_task
            fm_engine.dump_last_session(source="websocket_live")
            split_sessions.pop(session_id, None)
            print("[liveTryHeliumFM] websocket closed", flush=True)

    return app


# ---------------------------------------------------------------------------
# Entry point
# ---------------------------------------------------------------------------

def main() -> None:
    parser = LiveHeliumFMOptions()
    args = parser.parse()
    args.rank = args.device
    parser.print_options()

    app = build_app(args)

    import uvicorn

    print(f"[liveTryHeliumFM_ws_binary] serving {args.html_path} (binary av_transport)")
    print(f"[liveTryHeliumFM] open http://{args.host}:{args.port}/")
    uvicorn.run(app, host=args.host, port=args.port, log_level="info")


if __name__ == "__main__":
    main()