     1→"""
     2→Production inference engine for LeWM TTS.
     3→Features:
     4→  - Batched inference (many concurrent requests in parallel)
     5→  - KV-cache (O(1) per step instead of O(n))
     6→  - FP16
     7→  - Audio prompt support (for multi-speaker / voice cloning)
     8→  - Continuous batching (add/remove requests dynamically)
     9→  - Griffin-Lim + Vocos vocoder pipeline
    10→"""
    11→
    12→import torch
    13→import torch.nn.functional as F
    14→import torchaudio
    15→import numpy as np
    16→import time
    17→from dataclasses import dataclass, field
    18→from typing import Optional
    19→from pathlib import Path
    20→
    21→from model import LeWMTTS
    22→
    23→
    24→# ─── Vocoder ────────────────────────────────────────────────────────────────
    25→
    26→class Vocoder:
    27→    """Vocos vocoder for 100-mel → waveform (direct decode, no Griffin-Lim)."""
    28→
    29→    def __init__(self):
    30→        from vocos import Vocos
    31→        self.vocos = Vocos.from_pretrained("charactr/vocos-mel-24khz")
    32→        self.vocos.eval()
    33→
    34→    def __call__(self, mel):
    35→        """mel: [B, 100, T] → waveforms list of [T_audio] numpy arrays."""
    36→        with torch.no_grad():
    37→            # Vocos decode handles batched input
    38→            waveforms = []
    39→            for i in range(mel.shape[0]):
    40→                wav = self.vocos.decode(mel[i:i+1])  # [1, T_audio]
    41→                waveforms.append(wav.squeeze(0).numpy())
    42→        return waveforms
    43→
    44→
    45→# ─── Request / Result ───────────────────────────────────────────────────────
    46→
    47→@dataclass
    48→class TTSRequest:
    49→    text: str
    50→    max_steps: int = 300
    51→    temperature: float = 0.0
    52→    prompt_mel: Optional[torch.Tensor] = None  # [1, 80, T] for voice cloning
    53→
    54→
    55→@dataclass
    56→class TTSResult:
    57→    waveform: np.ndarray
    58→    sample_rate: int = 24000
    59→    steps_taken: int = 0
    60→    generation_time: float = 0.0
    61→
    62→    def save(self, path):
    63→        wav = torch.from_numpy(self.waveform).unsqueeze(0)
    64→        torchaudio.save(path, wav, self.sample_rate)
    65→
    66→    @property
    67→    def duration(self):
    68→        return len(self.waveform) / self.sample_rate
    69→
    70→
    71→# ─── Inference Engine ───────────────────────────────────────────────────────
    72→
    73→class InferenceEngine:
    74→    """
    75→    High-throughput batched TTS inference engine.
    76→
    77→    Usage:
    78→        engine = InferenceEngine("checkpoint.pt")
    79→
    80→        # Single request
    81→        result = engine.synthesize("नमस्ते भारत")
    82→
    83→        # Batched (high throughput)
    84→        results = engine.synthesize_batch(["text1", "text2", ...])
    85→
    86→        # With voice prompt
    87→        result = engine.synthesize("text", prompt_audio="ref.wav")
    88→
    89→        # Benchmark
    90→        engine.benchmark()
    91→    """
    92→
    93→    def __init__(self, checkpoint_path, device="cuda", dtype=torch.float16):
    94→        self.device = torch.device(device if torch.cuda.is_available() else "cpu")
    95→        self.dtype = dtype
    96→
    97→        # Load model
    98→        ckpt = torch.load(checkpoint_path, map_location=self.device, weights_only=False)
    99→        self.config = ckpt["config"]
   100→        self.model = LeWMTTS(self.config).to(self.device).to(dtype)
   101→        self.model.load_state_dict(ckpt["model"])
   102→        self.model.eval()
   103→
   104→        self.d_model = self.config["d_model"]
   105→
   106→        # Vocoder (runs on CPU)
   107→        self.vocoder = Vocoder()
   108→
   109→        # Precompute start embedding scale
   110→        self._start_scale = 5.0
   111→
   112→        print(f"Engine ready: {checkpoint_path}")
   113→        print(f"  Device: {self.device}, Dtype: {dtype}")
   114→        print(f"  Model: {sum(p.numel() for p in self.model.parameters())/1e6:.1f}M params")
   115→
   116→    # ─── Text encoding ──────────────────────────────────────────────────
   117→
   118→    def _encode_texts(self, texts):
   119→        """Encode list of texts → padded embeddings + mask.
   120→        Returns: text_emb [B, T_max, d], text_mask [B, T_max] (True=padding)
   121→        """
   122→        # Byte-level tokenization
   123→        token_lists = [list(t.encode("utf-8")) for t in texts]
   124→        max_len = max(len(t) for t in token_lists)
   125→
   126→        # Pad to same length
   127→        padded = torch.zeros(len(texts), max_len, dtype=torch.long, device=self.device)
   128→        mask = torch.ones(len(texts), max_len, dtype=torch.bool, device=self.device)
   129→
   130→        for i, tokens in enumerate(token_lists):
   131→            padded[i, :len(tokens)] = torch.tensor(tokens, dtype=torch.long)
   132→            mask[i, :len(tokens)] = False
   133→
   134→        with torch.no_grad():
   135→            text_emb = self.model.text_encoder(padded, text_mask=mask)
   136→
   137→        return text_emb, mask
   138→
   139→    # ─── Audio prompt encoding ──────────────────────────────────────────
   140→
   141→    def _encode_prompt(self, prompt_mel):
   142→        """Encode prompt mel → embeddings for AR seeding.
   143→        prompt_mel: [1, 80, T]
   144→        Returns: prompt_embs [1, T_down, d]
   145→        """
   146→        prompt_mel = prompt_mel.to(self.device).to(self.dtype)
   147→        with torch.no_grad():
   148→            embs = self.model.encode_audio(prompt_mel)  # [1, T_down, d]
   149→        return embs
   150→
   151→    def load_prompt_audio(self, audio_path, max_seconds=3.0):
   152→        """Load audio file → mel for prompting.
   153→        Returns: mel [1, 80, T]
   154→        """
   155→        wav, sr = torchaudio.load(audio_path)
   156→        if sr != 24000:
   157→            wav = torchaudio.functional.resample(wav, sr, 24000)
   158→        # Trim to max_seconds
   159→        max_samples = int(max_seconds * 24000)
   160→        wav = wav[:1, :max_samples]  # mono, trimmed
   161→
   162→        mel_transform = torchaudio.transforms.MelSpectrogram(
   163→            sample_rate=24000, n_fft=1024, hop_length=256,
   164→            n_mels=100, power=1.0,
   165→        )
   166→        mel = mel_transform(wav)
   167→        mel = torch.log(mel.clamp(min=1e-7))
   168→        return mel  # [1, 80, T]
   169→
   170→    # ─── Core AR generation (batched) ───────────────────────────────────
   171→
   172→    def _generate_batch(self, text_emb, text_mask, batch_size,
   173→                        max_steps=300, temperature=0.0,
   174→                        prompt_embs=None):
   175→        """
   176→        Core batched AR generation with KV-cache.
   177→
   178→        Args:
   179→            text_emb: [B, T_text, d]
   180→            text_mask: [B, T_text]
   181→            batch_size: int
   182→            max_steps: int
   183→            temperature: float
   184→            prompt_embs: list of [1, T_prompt, d] or None per request
   185→
   186→        Returns:
   187→            all_embs: [B, T_total, d]
   188→            steps_taken: list of int per request
   189→        """
   190→        d = self.d_model
   191→        cache = self.model.init_ar_cache()
   192→
   193→        # Track per-request state
   194→        active = torch.ones(batch_size, dtype=torch.bool, device=self.device)
   195→        steps_taken = [0] * batch_size
   196→        all_emb_lists = [[] for _ in range(batch_size)]
   197→        prev_norms = [[] for _ in range(batch_size)]
   198→
   199→        # Handle prompts: seed cache with prompt embeddings
   200→        if prompt_embs is not None:
   201→            max_prompt_len = max(
   202→                (e.shape[1] for e in prompt_embs if e is not None), default=0
   203→            )
   204→            if max_prompt_len > 0:
   205→                # Process prompt embeddings through cache step by step
   206→                # Pad prompts to same length
   207→                prompt_padded = torch.zeros(
   208→                    batch_size, max_prompt_len, d,
   209→                    device=self.device, dtype=self.dtype,
   210→                )
   211→                prompt_lengths = []
   212→                for i, pe in enumerate(prompt_embs):
   213→                    if pe is not None:
   214→                        L = pe.shape[1]
   215→                        prompt_padded[i, :L] = pe[0]
   216→                        prompt_lengths.append(L)
   217→                    else:
   218→                        prompt_lengths.append(0)
   219→
   220→                # Feed prompt through cache one step at a time
   221→                for t in range(max_prompt_len):
   222→                    step_emb = prompt_padded[:, t:t+1, :]  # [B, 1, d]
   223→                    _, cache = self.model.predict_next_cached(
   224→                        step_emb, text_emb, t, cache, text_mask=text_mask,
   225→                    )
   226→                    # Store embeddings
   227→                    for i in range(batch_size):
   228→                        if t < prompt_lengths[i]:
   229→                            all_emb_lists[i].append(step_emb[i:i+1])
   230→
   231→                step_offset = max_prompt_len
   232→                # Use last prompt prediction as first AR input
   233→                last_pred, cache = self.model.predict_next_cached(
   234→                    prompt_padded[:, -1:, :], text_emb, step_offset, cache,
   235→                    text_mask=text_mask,
   236→                )
   237→                current_emb = last_pred
   238→                step_offset += 1
   239→            else:
   240→                step_offset = 0
   241→                current_emb = None
   242→        else:
   243→            step_offset = 0
   244→            current_emb = None
   245→
   246→        # If no prompt, start with random embedding
   247→        if current_emb is None:
   248→            start_emb = torch.randn(
   249→                batch_size, 1, d, device=self.device, dtype=self.dtype,
   250→            ) * self._start_scale
   251→            current_emb, cache = self.model.predict_next_cached(
   252→                start_emb, text_emb, 0, cache, text_mask=text_mask,
   253→            )
   254→            for i in range(batch_size):
   255→                all_emb_lists[i].append(start_emb[i:i+1])
   256→            step_offset = 1
   257→
   258→        if temperature > 0:
   259→            current_emb = current_emb + torch.randn_like(current_emb) * temperature
   260→
   261→        for i in range(batch_size):
   262→            all_emb_lists[i].append(current_emb[i:i+1])
   263→
   264→        # AR loop
   265→        for step in range(1, max_steps):
   266→            if not active.any():
   267→                break
   268→
   269→            next_emb, cache = self.model.predict_next_cached(
   270→                current_emb, text_emb, step + step_offset, cache,
   271→                text_mask=text_mask,
   272→            )
   273→
   274→            if temperature > 0:
   275→                next_emb = next_emb + torch.randn_like(next_emb) * temperature
   276→
   277→            # Store and check stopping per request
   278→            for i in range(batch_size):
   279→                if not active[i]:
   280→                    continue
   281→
   282→                all_emb_lists[i].append(next_emb[i:i+1])
   283→                steps_taken[i] = step
   284→
   285→                norm = next_emb[i].norm().item()
   286→                prev_norms[i].append(norm)
   287→                if len(prev_norms[i]) > 20:
   288→                    recent = prev_norms[i][-20:]
   289→                    if np.std(recent) < 0.01 * np.mean(recent) and step > 30:
   290→                        active[i] = False
   291→
   292→            current_emb = next_emb
   293→
   294→        # Concatenate embeddings per request
   295→        results = []
   296→        for i in range(batch_size):
   297→            embs = torch.cat(all_emb_lists[i], dim=1)  # [1, T, d]
   298→            results.append(embs)
   299→
   300→        return results, steps_taken
   301→
   302→    # ─── Public API ─────────────────────────────────────────────────────
   303→
   304→    def synthesize(self, text, max_steps=300, temperature=0.0,
   305→                   prompt_audio=None, prompt_seconds=3.0):
   306→        """Synthesize single text.
   307→
   308→        Args:
   309→            text: Hindi text string
   310→            max_steps: max AR steps
   311→            temperature: sampling temperature (0 = greedy)
   312→            prompt_audio: path to reference audio for voice cloning
   313→            prompt_seconds: max seconds of prompt to use
   314→
   315→        Returns: TTSResult
   316→        """
   317→        return self.synthesize_batch(
   318→            [text], max_steps=max_steps, temperature=temperature,
   319→            prompt_audios=[prompt_audio] if prompt_audio else None,
   320→            prompt_seconds=prompt_seconds,
   321→        )[0]
   322→
   323→    def synthesize_batch(self, texts, max_steps=300, temperature=0.0,
   324→                         prompt_audios=None, prompt_seconds=3.0):
   325→        """Synthesize a batch of texts in parallel.
   326→
   327→        Args:
   328→            texts: list of text strings
   329→            max_steps: max AR steps
   330→            temperature: sampling temperature
   331→            prompt_audios: list of audio paths (or None) per text
   332→            prompt_seconds: max seconds of prompt
   333→
   334→        Returns: list of TTSResult
   335→        """
   336→        B = len(texts)
   337→        t0 = time.time()
   338→
   339→        # Encode texts
   340→        text_emb, text_mask = self._encode_texts(texts)
   341→
   342→        # Encode prompts if given
   343→        prompt_embs = None
   344→        if prompt_audios is not None:
   345→            prompt_embs = []
   346→            for pa in prompt_audios:
   347→                if pa is not None:
   348→                    mel = self.load_prompt_audio(pa, max_seconds=prompt_seconds)
   349→                    embs = self._encode_prompt(mel)
   350→                    prompt_embs.append(embs)
   351→                else:
   352→                    prompt_embs.append(None)
   353→
   354→        # Generate
   355→        with torch.no_grad():
   356→            emb_results, steps_taken = self._generate_batch(
   357→                text_emb, text_mask, B,
   358→                max_steps=max_steps, temperature=temperature,
   359→                prompt_embs=prompt_embs,
   360→            )
   361→
   362→        # Decode to mel → waveform
   363→        results = []
   364→        for i in range(B):
   365→            with torch.no_grad():
   366→                mel = self.model.mel_decoder(emb_results[i].to(self.device))
   367→                # mel: [1, 80, T]
   368→            wav_list = self.vocoder(mel.float().cpu())
   369→            waveform = wav_list[0]
   370→
   371→            # Normalize
   372→            mx = np.abs(waveform).max()
   373→            if mx > 0:
   374→                waveform = waveform / mx * 0.95
   375→
   376→            results.append(TTSResult(
   377→                waveform=waveform,
   378→                steps_taken=steps_taken[i],
   379→                generation_time=time.time() - t0,
   380→            ))
   381→
   382→        return results
   383→
   384→    # ─── Benchmark ──────────────────────────────────────────────────────
   385→
   386→    def benchmark(self, batch_sizes=None, steps=300):
   387→        """Benchmark throughput at various batch sizes."""
   388→        if batch_sizes is None:
   389→            batch_sizes = [1, 4, 16, 32, 64, 128, 256]
   390→
   391→        text = "भारत एक महान देश है और हम सब भारतवासी हैं।"
   392→        audio_per_step = 4 * 256 / 24000  # seconds of audio per AR step
   393→
   394→        print(f"\n{'='*70}")
   395→        print(f"  LeWM TTS Inference Benchmark — {self.device}, {self.dtype}")
   396→        print(f"  Model: {self.config['d_model']}d, {self.config['predictor_layers']}L predictor")
   397→        print(f"  AR steps: {steps}")
   398→        print(f"{'='*70}")
   399→        print(f"  {'Batch':>6} {'Time(s)':>8} {'Audio(s)':>9} {'RTF':>8} "
   400→              f"{'Speed':>8} {'Throughput':>12} {'Per-req':>10}")
   401→        print(f"  {'-'*6} {'-'*8} {'-'*9} {'-'*8} {'-'*8} {'-'*12} {'-'*10}")
   402→
   403→        for B in batch_sizes:
   404→            try:
   405→                texts = [text] * B
   406→                token_lists = [list(t.encode("utf-8")) for t in texts]
   407→                max_len = max(len(t) for t in token_lists)
   408→                padded = torch.zeros(B, max_len, dtype=torch.long, device=self.device)
   409→                mask = torch.ones(B, max_len, dtype=torch.bool, device=self.device)
   410→                for i, tokens in enumerate(token_lists):
   411→                    padded[i, :len(tokens)] = torch.tensor(tokens, dtype=torch.long)
   412→                    mask[i, :len(tokens)] = False
   413→
   414→                with torch.no_grad():
   415→                    text_emb = self.model.text_encoder(padded, text_mask=mask)
   416→
   417→                # Warmup
   418→                if B == batch_sizes[0]:
   419→                    cache = self.model.init_ar_cache()
   420→                    start = torch.randn(B, 1, self.d_model,
   421→                                        device=self.device, dtype=self.dtype) * self._start_scale
   422→                    nxt = start
   423→                    for s in range(5):
   424→                        nxt, cache = self.model.predict_next_cached(
   425→                            nxt, text_emb, s, cache, text_mask=mask)
   426→                    torch.cuda.synchronize()
   427→
   428→                # Benchmark
   429→                cache = self.model.init_ar_cache()
   430→                start = torch.randn(B, 1, self.d_model,
   431→                                    device=self.device, dtype=self.dtype) * self._start_scale
   432→                nxt = start
   433→
   434→                torch.cuda.synchronize()
   435→                t0 = time.time()
   436→                for s in range(steps):
   437→                    nxt, cache = self.model.predict_next_cached(
   438→                        nxt, text_emb, s, cache, text_mask=mask)
   439→                torch.cuda.synchronize()
   440→                elapsed = time.time() - t0
   441→
   442→                total_audio = B * steps * audio_per_step
   443→                rtf = elapsed / (steps * audio_per_step)
   444→                speed = (steps * audio_per_step) / elapsed
   445→                throughput = total_audio / elapsed
   446→                per_req = elapsed / B
   447→
   448→                print(f"  {B:>6} {elapsed:>8.3f} {total_audio:>9.1f} {rtf:>8.4f} "
   449→                      f"{speed:>7.0f}x {throughput:>10.0f}x {per_req:>9.3f}s")
   450→
   451→            except torch.cuda.OutOfMemoryError:
   452→                print(f"  {B:>6} {'OOM':>8}")
   453→                torch.cuda.empty_cache()
   454→                break
   455→
   456→        print(f"{'='*70}")
   457→        print(f"  Throughput = total audio-seconds generated per wall-second")
   458→        print(f"  Speed = single-request real-time factor")
   459→        print(f"  Per-req = wall-time per request in batch")
   460→        print()
   461→
   462→
   463→# ─── CLI ────────────────────────────────────────────────────────────────────
   464→
   465→if __name__ == "__main__":
   466→    import argparse
   467→
   468→    parser = argparse.ArgumentParser(description="LeWM TTS Inference Engine")
   469→    parser.add_argument("--checkpoint", required=True)
   470→    parser.add_argument("--text", default=None, help="Text to synthesize")
   471→    parser.add_argument("--texts", nargs="+", help="Multiple texts for batch synthesis")
   472→    parser.add_argument("--prompt", default=None, help="Reference audio for voice cloning")
   473→    parser.add_argument("--output", default="output.wav")
   474→    parser.add_argument("--output_dir", default=None, help="Output dir for batch mode")
   475→    parser.add_argument("--max_steps", type=int, default=300)
   476→    parser.add_argument("--temperature", type=float, default=0.0)
   477→    parser.add_argument("--benchmark", action="store_true", help="Run throughput benchmark")
   478→    parser.add_argument("--device", default="cuda")
   479→    parser.add_argument("--fp32", action="store_true", help="Use FP32 instead of FP16")
   480→    args = parser.parse_args()
   481→
   482→    dtype = torch.float32 if args.fp32 else torch.float16
   483→    engine = InferenceEngine(args.checkpoint, device=args.device, dtype=dtype)
   484→
   485→    if args.benchmark:
   486→        engine.benchmark()
   487→
   488→    elif args.texts:
   489→        results = engine.synthesize_batch(
   490→            args.texts, max_steps=args.max_steps, temperature=args.temperature,
   491→            prompt_audios=[args.prompt] * len(args.texts) if args.prompt else None,
   492→        )
   493→        out_dir = Path(args.output_dir or "output")
   494→        out_dir.mkdir(exist_ok=True)
   495→        for i, r in enumerate(results):
   496→            path = out_dir / f"batch_{i:03d}.wav"
   497→            r.save(str(path))
   498→            print(f"  [{i}] {path} — {r.duration:.2f}s, {r.steps_taken} steps")
   499→
   500→    elif args.text:
   501→        result = engine.synthesize(
   502→            args.text, max_steps=args.max_steps, temperature=args.temperature,
   503→            prompt_audio=args.prompt,
   504→        )
   505→        result.save(args.output)
   506→        print(f"Saved: {args.output} — {result.duration:.2f}s, "
   507→              f"{result.steps_taken} steps, {result.generation_time:.2f}s")
   508→    else:
   509→        parser.print_help()
   510→

"""
Production inference engine for LeWM TTS.
Features:
  - Batched inference (many concurrent requests in parallel)
  - KV-cache (O(1) per step instead of O(n))
  - FP16
  - Audio prompt support (for multi-speaker / voice cloning)
  - Continuous batching (add/remove requests dynamically)
  - Griffin-Lim + Vocos vocoder pipeline
"""

import torch
import torch.nn.functional as F
import torchaudio
import numpy as np
import time
from dataclasses import dataclass, field
from typing import Optional
from pathlib import Path

from model import LeWMTTS

# ─── Vocoder ────────────────────────────────────────────────────────────────

class Vocoder:
    """Vocos vocoder for 100-mel → waveform (direct decode, no Griffin-Lim)."""

    def __init__(self):
        from vocos import Vocos
        self.vocos = Vocos.from_pretrained("charactr/vocos-mel-24khz")
        self.vocos.eval()

    def __call__(self, mel):
        """mel: [B, 100, T] → waveforms list of [T_audio] numpy arrays."""
        with torch.no_grad():
            # Vocos decode handles batched input
            waveforms = []
            for i in range(mel.shape[0]):
                wav = self.vocos.decode(mel[i:i+1])  # [1, T_audio]
                waveforms.append(wav.squeeze(0).numpy())
        return waveforms

# ─── Request / Result ───────────────────────────────────────────────────────

@dataclass
class TTSRequest:
    text: str
    max_steps: int = 300
    temperature: float = 0.0
    prompt_mel: Optional[torch.Tensor] = None  # [1, 80, T] for voice cloning

@dataclass
class TTSResult:
    waveform: np.ndarray
    sample_rate: int = 24000
    steps_taken: int = 0
    generation_time: float = 0.0

    def save(self, path):
        wav = torch.from_numpy(self.waveform).unsqueeze(0)
        torchaudio.save(path, wav, self.sample_rate)

    @property
    def duration(self):
        return len(self.waveform) / self.sample_rate

# ─── Inference Engine ───────────────────────────────────────────────────────

class InferenceEngine:
    """
    High-throughput batched TTS inference engine.

    Usage:
        engine = InferenceEngine("checkpoint.pt")

        # Single request
        result = engine.synthesize("नमस्ते भारत")

        # Batched (high throughput)
        results = engine.synthesize_batch(["text1", "text2", ...])

        # With voice prompt
        result = engine.synthesize("text", prompt_audio="ref.wav")

        # Benchmark
        engine.benchmark()
    """

    def __init__(self, checkpoint_path, device="cuda", dtype=torch.float16):
        self.device = torch.device(device if torch.cuda.is_available() else "cpu")
        self.dtype = dtype

        # Load model
        ckpt = torch.load(checkpoint_path, map_location=self.device, weights_only=False)
        self.config = ckpt["config"]
        self.model = LeWMTTS(self.config).to(self.device).to(dtype)
        self.model.load_state_dict(ckpt["model"])
        self.model.eval()

        self.d_model = self.config["d_model"]

        # Vocoder (runs on CPU)
        self.vocoder = Vocoder()

        # Precompute start embedding scale
        self._start_scale = 5.0

        print(f"Engine ready: {checkpoint_path}")
        print(f"  Device: {self.device}, Dtype: {dtype}")
        print(f"  Model: {sum(p.numel() for p in self.model.parameters())/1e6:.1f}M params")

    # ─── Text encoding ──────────────────────────────────────────────────

    def _encode_texts(self, texts):
        """Encode list of texts → padded embeddings + mask.
        Returns: text_emb [B, T_max, d], text_mask [B, T_max] (True=padding)
        """
        # Byte-level tokenization
        token_lists = [list(t.encode("utf-8")) for t in texts]
        max_len = max(len(t) for t in token_lists)

        # Pad to same length
        padded = torch.zeros(len(texts), max_len, dtype=torch.long, device=self.device)
        mask = torch.ones(len(texts), max_len, dtype=torch.bool, device=self.device)

        for i, tokens in enumerate(token_lists):
            padded[i, :len(tokens)] = torch.tensor(tokens, dtype=torch.long)
            mask[i, :len(tokens)] = False

        with torch.no_grad():
            text_emb = self.model.text_encoder(padded, text_mask=mask)

        return text_emb, mask

    # ─── Audio prompt encoding ──────────────────────────────────────────

    def _encode_prompt(self, prompt_mel):
        """Encode prompt mel → embeddings for AR seeding.
        prompt_mel: [1, 80, T]
        Returns: prompt_embs [1, T_down, d]
        """
        prompt_mel = prompt_mel.to(self.device).to(self.dtype)
        with torch.no_grad():
            embs = self.model.encode_audio(prompt_mel)  # [1, T_down, d]
        return embs

    def load_prompt_audio(self, audio_path, max_seconds=3.0):
        """Load audio file → mel for prompting.
        Returns: mel [1, 80, T]
        """
        wav, sr = torchaudio.load(audio_path)
        if sr != 24000:
            wav = torchaudio.functional.resample(wav, sr, 24000)
        # Trim to max_seconds
        max_samples = int(max_seconds * 24000)
        wav = wav[:1, :max_samples]  # mono, trimmed

        mel_transform = torchaudio.transforms.MelSpectrogram(
            sample_rate=24000, n_fft=1024, hop_length=256,
            n_mels=100, power=1.0,
        )
        mel = mel_transform(wav)
        mel = torch.log(mel.clamp(min=1e-7))
        return mel  # [1, 80, T]

    # ─── Core AR generation (batched) ───────────────────────────────────

    def _generate_batch(self, text_emb, text_mask, batch_size,
                        max_steps=300, temperature=0.0,
                        prompt_embs=None):
        """
        Core batched AR generation with KV-cache.

        Args:
            text_emb: [B, T_text, d]
            text_mask: [B, T_text]
            batch_size: int
            max_steps: int
            temperature: float
            prompt_embs: list of [1, T_prompt, d] or None per request

        Returns:
            all_embs: [B, T_total, d]
            steps_taken: list of int per request
        """
        d = self.d_model
        cache = self.model.init_ar_cache()

        # Track per-request state
        active = torch.ones(batch_size, dtype=torch.bool, device=self.device)
        steps_taken = [0] * batch_size
        all_emb_lists = [[] for _ in range(batch_size)]
        prev_norms = [[] for _ in range(batch_size)]

        # Handle prompts: seed cache with prompt embeddings
        if prompt_embs is not None:
            max_prompt_len = max(
                (e.shape[1] for e in prompt_embs if e is not None), default=0
            )
            if max_prompt_len > 0:
                # Process prompt embeddings through cache step by step
                # Pad prompts to same length
                prompt_padded = torch.zeros(
                    batch_size, max_prompt_len, d,
                    device=self.device, dtype=self.dtype,
                )
                prompt_lengths = []
                for i, pe in enumerate(prompt_embs):
                    if pe is not None:
                        L = pe.shape[1]
                        prompt_padded[i, :L] = pe[0]
                        prompt_lengths.append(L)
                    else:
                        prompt_lengths.append(0)

                # Feed prompt through cache one step at a time
                for t in range(max_prompt_len):
                    step_emb = prompt_padded[:, t:t+1, :]  # [B, 1, d]
                    _, cache = self.model.predict_next_cached(
                        step_emb, text_emb, t, cache, text_mask=text_mask,
                    )
                    # Store embeddings
                    for i in range(batch_size):
                        if t < prompt_lengths[i]:
                            all_emb_lists[i].append(step_emb[i:i+1])

                step_offset = max_prompt_len
                # Use last prompt prediction as first AR input
                last_pred, cache = self.model.predict_next_cached(
                    prompt_padded[:, -1:, :], text_emb, step_offset, cache,
                    text_mask=text_mask,
                )
                current_emb = last_pred
                step_offset += 1
            else:
                step_offset = 0
                current_emb = None
        else:
            step_offset = 0
            current_emb = None

        # If no prompt, start with random embedding
        if current_emb is None:
            start_emb = torch.randn(
                batch_size, 1, d, device=self.device, dtype=self.dtype,
            ) * self._start_scale
            current_emb, cache = self.model.predict_next_cached(
                start_emb, text_emb, 0, cache, text_mask=text_mask,
            )
            for i in range(batch_size):
                all_emb_lists[i].append(start_emb[i:i+1])
            step_offset = 1

        if temperature > 0:
            current_emb = current_emb + torch.randn_like(current_emb) * temperature

        for i in range(batch_size):
            all_emb_lists[i].append(current_emb[i:i+1])

        # AR loop
        for step in range(1, max_steps):
            if not active.any():
                break

            next_emb, cache = self.model.predict_next_cached(
                current_emb, text_emb, step + step_offset, cache,
                text_mask=text_mask,
            )

            if temperature > 0:
                next_emb = next_emb + torch.randn_like(next_emb) * temperature

            # Store and check stopping per request
            for i in range(batch_size):
                if not active[i]:
                    continue

                all_emb_lists[i].append(next_emb[i:i+1])
                steps_taken[i] = step

                norm = next_emb[i].norm().item()
                prev_norms[i].append(norm)
                if len(prev_norms[i]) > 20:
                    recent = prev_norms[i][-20:]
                    if np.std(recent) < 0.01 * np.mean(recent) and step > 30:
                        active[i] = False

            current_emb = next_emb

        # Concatenate embeddings per request
        results = []
        for i in range(batch_size):
            embs = torch.cat(all_emb_lists[i], dim=1)  # [1, T, d]
            results.append(embs)

        return results, steps_taken

    # ─── Public API ─────────────────────────────────────────────────────

    def synthesize(self, text, max_steps=300, temperature=0.0,
                   prompt_audio=None, prompt_seconds=3.0):
        """Synthesize single text.

        Args:
            text: Hindi text string
            max_steps: max AR steps
            temperature: sampling temperature (0 = greedy)
            prompt_audio: path to reference audio for voice cloning
            prompt_seconds: max seconds of prompt to use

        Returns: TTSResult
        """
        return self.synthesize_batch(
            [text], max_steps=max_steps, temperature=temperature,
            prompt_audios=[prompt_audio] if prompt_audio else None,
            prompt_seconds=prompt_seconds,
        )[0]

    def synthesize_batch(self, texts, max_steps=300, temperature=0.0,
                         prompt_audios=None, prompt_seconds=3.0):
        """Synthesize a batch of texts in parallel.

        Args:
            texts: list of text strings
            max_steps: max AR steps
            temperature: sampling temperature
            prompt_audios: list of audio paths (or None) per text
            prompt_seconds: max seconds of prompt

        Returns: list of TTSResult
        """
        B = len(texts)
        t0 = time.time()

        # Encode texts
        text_emb, text_mask = self._encode_texts(texts)

        # Encode prompts if given
        prompt_embs = None
        if prompt_audios is not None:
            prompt_embs = []
            for pa in prompt_audios:
                if pa is not None:
                    mel = self.load_prompt_audio(pa, max_seconds=prompt_seconds)
                    embs = self._encode_prompt(mel)
                    prompt_embs.append(embs)
                else:
                    prompt_embs.append(None)

        # Generate
        with torch.no_grad():
            emb_results, steps_taken = self._generate_batch(
                text_emb, text_mask, B,
                max_steps=max_steps, temperature=temperature,
                prompt_embs=prompt_embs,
            )

        # Decode to mel → waveform
        results = []
        for i in range(B):
            with torch.no_grad():
                mel = self.model.mel_decoder(emb_results[i].to(self.device))
                # mel: [1, 80, T]
            wav_list = self.vocoder(mel.float().cpu())
            waveform = wav_list[0]

            # Normalize
            mx = np.abs(waveform).max()
            if mx > 0:
                waveform = waveform / mx * 0.95

            results.append(TTSResult(
                waveform=waveform,
                steps_taken=steps_taken[i],
                generation_time=time.time() - t0,
            ))

        return results

    # ─── Benchmark ──────────────────────────────────────────────────────

    def benchmark(self, batch_sizes=None, steps=300):
        """Benchmark throughput at various batch sizes."""
        if batch_sizes is None:
            batch_sizes = [1, 4, 16, 32, 64, 128, 256]

        text = "भारत एक महान देश है और हम सब भारतवासी हैं।"
        audio_per_step = 4 * 256 / 24000  # seconds of audio per AR step

        print(f"\n{'='*70}")
        print(f"  LeWM TTS Inference Benchmark — {self.device}, {self.dtype}")
        print(f"  Model: {self.config['d_model']}d, {self.config['predictor_layers']}L predictor")
        print(f"  AR steps: {steps}")
        print(f"{'='*70}")
        print(f"  {'Batch':>6} {'Time(s)':>8} {'Audio(s)':>9} {'RTF':>8} "
              f"{'Speed':>8} {'Throughput':>12} {'Per-req':>10}")
        print(f"  {'-'*6} {'-'*8} {'-'*9} {'-'*8} {'-'*8} {'-'*12} {'-'*10}")

        for B in batch_sizes:
            try:
                texts = [text] * B
                token_lists = [list(t.encode("utf-8")) for t in texts]
                max_len = max(len(t) for t in token_lists)
                padded = torch.zeros(B, max_len, dtype=torch.long, device=self.device)
                mask = torch.ones(B, max_len, dtype=torch.bool, device=self.device)
                for i, tokens in enumerate(token_lists):
                    padded[i, :len(tokens)] = torch.tensor(tokens, dtype=torch.long)
                    mask[i, :len(tokens)] = False

                with torch.no_grad():
                    text_emb = self.model.text_encoder(padded, text_mask=mask)

                # Warmup
                if B == batch_sizes[0]:
                    cache = self.model.init_ar_cache()
                    start = torch.randn(B, 1, self.d_model,
                                        device=self.device, dtype=self.dtype) * self._start_scale
                    nxt = start
                    for s in range(5):
                        nxt, cache = self.model.predict_next_cached(
                            nxt, text_emb, s, cache, text_mask=mask)
                    torch.cuda.synchronize()

                # Benchmark
                cache = self.model.init_ar_cache()
                start = torch.randn(B, 1, self.d_model,
                                    device=self.device, dtype=self.dtype) * self._start_scale
                nxt = start

                torch.cuda.synchronize()
                t0 = time.time()
                for s in range(steps):
                    nxt, cache = self.model.predict_next_cached(
                        nxt, text_emb, s, cache, text_mask=mask)
                torch.cuda.synchronize()
                elapsed = time.time() - t0

                total_audio = B * steps * audio_per_step
                rtf = elapsed / (steps * audio_per_step)
                speed = (steps * audio_per_step) / elapsed
                throughput = total_audio / elapsed
                per_req = elapsed / B

                print(f"  {B:>6} {elapsed:>8.3f} {total_audio:>9.1f} {rtf:>8.4f} "
                      f"{speed:>7.0f}x {throughput:>10.0f}x {per_req:>9.3f}s")

            except torch.cuda.OutOfMemoryError:
                print(f"  {B:>6} {'OOM':>8}")
                torch.cuda.empty_cache()
                break

        print(f"{'='*70}")
        print(f"  Throughput = total audio-seconds generated per wall-second")
        print(f"  Speed = single-request real-time factor")
        print(f"  Per-req = wall-time per request in batch")
        print()

# ─── CLI ────────────────────────────────────────────────────────────────────

if __name__ == "__main__":
    import argparse

    parser = argparse.ArgumentParser(description="LeWM TTS Inference Engine")
    parser.add_argument("--checkpoint", required=True)
    parser.add_argument("--text", default=None, help="Text to synthesize")
    parser.add_argument("--texts", nargs="+", help="Multiple texts for batch synthesis")
    parser.add_argument("--prompt", default=None, help="Reference audio for voice cloning")
    parser.add_argument("--output", default="output.wav")
    parser.add_argument("--output_dir", default=None, help="Output dir for batch mode")
    parser.add_argument("--max_steps", type=int, default=300)
    parser.add_argument("--temperature", type=float, default=0.0)
    parser.add_argument("--benchmark", action="store_true", help="Run throughput benchmark")
    parser.add_argument("--device", default="cuda")
    parser.add_argument("--fp32", action="store_true", help="Use FP32 instead of FP16")
    args = parser.parse_args()

    dtype = torch.float32 if args.fp32 else torch.float16
    engine = InferenceEngine(args.checkpoint, device=args.device, dtype=dtype)

    if args.benchmark:
        engine.benchmark()

    elif args.texts:
        results = engine.synthesize_batch(
            args.texts, max_steps=args.max_steps, temperature=args.temperature,
            prompt_audios=[args.prompt] * len(args.texts) if args.prompt else None,
        )
        out_dir = Path(args.output_dir or "output")
        out_dir.mkdir(exist_ok=True)
        for i, r in enumerate(results):
            path = out_dir / f"batch_{i:03d}.wav"
            r.save(str(path))
            print(f"  [{i}] {path} — {r.duration:.2f}s, {r.steps_taken} steps")

    elif args.text:
        result = engine.synthesize(
            args.text, max_steps=args.max_steps, temperature=args.temperature,
            prompt_audio=args.prompt,
        )
        result.save(args.output)
        print(f"Saved: {args.output} — {result.duration:.2f}s, "
              f"{result.steps_taken} steps, {result.generation_time:.2f}s")
    else:
        parser.print_help()


Found more issues:

1. **inference_engine.py lines 247-256**: Uses `torch.randn * self._start_scale` (random start embedding). Need to update to use `self.model.start_emb`.

2. **inference_engine.py line 110**: `self._start_scale = 5.0` — should be removed.

3. **inference_engine.py benchmark**: Also uses random start embeddings (lines 420-421, 430-431).

4. **Model.py forward — EMA encoder dropout during training**: The EMA encoder runs in training mode (dropout active) because `self.ema_audio_encoder` inherits the model's training state.

I need to set the EMA encoder to eval mode during the forward pass so it produces deterministic targets, then fix all these issues.


Found two more issues:

**Issue 1**: `inference_engine.py` still uses random start embeddings (`randn * 5.0`) in multiple places.

**Issue 2**: The EMA encoder runs with **dropout enabled** during training (it inherits the model's `.training` state). EMA targets should be deterministic — need to force eval mode.

Let me fix both:


Edit

        # EMA target encoder → stable prediction targets (no gradient)
        with torch.no_grad():
            _, mu_ema, _ = self.ema_audio_encoder(mel, mel_mask)
            target_emb = mu_ema  # [B, T_down, d] — all EMA targets

        # EMA target encoder → stable prediction targets (no gradient)
        # Force eval mode so dropout is disabled — EMA targets must be deterministic
        with torch.no_grad():
            self.ema_audio_encoder.eval()
            _, mu_ema, _ = self.ema_audio_encoder(mel, mel_mask)
            self.ema_audio_encoder.train()
            target_emb = mu_ema  # [B, T_down, d] — all EMA targets


The file /home/ubuntu/lewm-tts/model.py has been updated successfully.

        # EMA target encoder → stable prediction targets (no gradient)
        # Force eval mode so dropout is disabled — EMA targets must be deterministic
        with torch.no_grad():
            self.ema_audio_encoder.eval()
            _, mu_ema, _ = self.ema_audio_encoder(mel, mel_mask)
            self.ema_audio_encoder.train()
            target_emb = mu_ema  # [B, T_down, d] — all EMA targets

"""
LeWM TTS Model — JEPA-based Text-to-Speech

Architecture:
  - TextEncoder: Character-level transformer for Hindi/Devanagari text
  - AudioEncoder: 1D CNN + Transformer that encodes mel spectrograms to embeddings
  - JEPAPredictor: Transformer that predicts next audio embedding given text + audio context
  - Two losses: next-embedding prediction (MSE) + Gaussian regularizer (KL to N(0,I))
"""

import math
import copy
import torch
import torch.nn as nn
import torch.nn.functional as F
import torchaudio

# ─── Text Encoder ───────────────────────────────────────────────────────────

class TextEncoder(nn.Module):
    """Character-level transformer encoder for Hindi text."""

    def __init__(self, vocab_size=256, d_model=256, nhead=4, num_layers=4, dropout=0.1):
        super().__init__()
        self.d_model = d_model

        # Character embedding (covers all Unicode Hindi chars via byte-level encoding)
        self.char_embed = nn.Embedding(vocab_size, d_model)
        self.pos_embed = SinusoidalPositionalEncoding(d_model, max_len=2048)
        self.dropout = nn.Dropout(dropout)

        encoder_layer = nn.TransformerEncoderLayer(
            d_model=d_model, nhead=nhead,
            dim_feedforward=d_model * 4,
            dropout=dropout, batch_first=True,
            activation="gelu",
        )
        self.transformer = nn.TransformerEncoder(encoder_layer, num_layers=num_layers)
        self.proj = nn.Linear(d_model, d_model)

    def forward(self, text_tokens, text_mask=None):
        """
        Args:
            text_tokens: [B, T_text] long tensor of byte-level token IDs
            text_mask: [B, T_text] bool, True = padding
        Returns:
            text_emb: [B, T_text, d_model]
        """
        x = self.char_embed(text_tokens) * math.sqrt(self.d_model)
        x = self.pos_embed(x)
        x = self.dropout(x)
        x = self.transformer(x, src_key_padding_mask=text_mask)
        x = self.proj(x)
        return x

# ─── Audio Encoder ──────────────────────────────────────────────────────────

class AudioEncoder(nn.Module):
    """
    Encodes mel spectrograms into a sequence of audio embeddings.
    Uses 1D CNN for downsampling + Transformer for contextualization.
    """

    def __init__(self, n_mels=100, d_model=256, nhead=4, num_layers=4,
                 downsample_factor=4, dropout=0.1):
        super().__init__()
        self.d_model = d_model
        self.downsample_factor = downsample_factor

        # CNN downsampling: reduces temporal resolution by downsample_factor
        if downsample_factor == 2:
            self.conv_pre = nn.Sequential(
                nn.Conv1d(n_mels, d_model, kernel_size=7, padding=3),
                nn.GELU(),
                nn.Conv1d(d_model, d_model, kernel_size=4, stride=2, padding=1),
                nn.GELU(),
            )  # Total stride = 2
        else:
            self.conv_pre = nn.Sequential(
                nn.Conv1d(n_mels, d_model, kernel_size=7, padding=3),
                nn.GELU(),
                nn.Conv1d(d_model, d_model, kernel_size=4, stride=2, padding=1),
                nn.GELU(),
                nn.Conv1d(d_model, d_model, kernel_size=4, stride=2, padding=1),
                nn.GELU(),
            )  # Total stride = 4

        self.pos_embed = SinusoidalPositionalEncoding(d_model, max_len=4096)
        self.dropout = nn.Dropout(dropout)

        encoder_layer = nn.TransformerEncoderLayer(
            d_model=d_model, nhead=nhead,
            dim_feedforward=d_model * 4,
            dropout=dropout, batch_first=True,
            activation="gelu",
        )
        self.transformer = nn.TransformerEncoder(encoder_layer, num_layers=num_layers)

        # Project to embedding with mu and logvar for Gaussian regularizer
        self.proj_mu = nn.Linear(d_model, d_model)
        self.proj_logvar = nn.Linear(d_model, d_model)

    def forward(self, mel, mel_mask=None):
        """
        Args:
            mel: [B, n_mels, T_mel] — log mel spectrogram
            mel_mask: [B, T_mel] bool, True = padding (before downsampling)
        Returns:
            z: [B, T_down, d_model] — sampled latent embeddings
            mu: [B, T_down, d_model]
            logvar: [B, T_down, d_model]
        """
        # CNN encoding: [B, n_mels, T] -> [B, d_model, T//4]
        x = self.conv_pre(mel)
        x = x.transpose(1, 2)  # [B, T_down, d_model]

        # Downsample mask if provided
        if mel_mask is not None:
            # Approximate downsampled mask
            T_down = x.shape[1]
            mel_mask = mel_mask[:, ::self.downsample_factor][:, :T_down]

        x = self.pos_embed(x)
        x = self.dropout(x)
        x = self.transformer(x, src_key_padding_mask=mel_mask)

        mu = self.proj_mu(x)
        logvar = self.proj_logvar(x)

        # Reparameterization trick
        if self.training:
            std = torch.exp(0.5 * logvar)
            eps = torch.randn_like(std)
            z = mu + eps * std
        else:
            z = mu

        return z, mu, logvar

# ─── JEPA Predictor ─────────────────────────────────────────────────────────

class JEPAPredictor(nn.Module):
    """
    Predicts next audio embedding given:
    - Text context (from TextEncoder)
    - Previous audio embeddings (from AudioEncoder)

    Uses cross-attention to condition on text, and causal self-attention
    over the audio embedding sequence.
    """

    def __init__(self, d_model=256, nhead=4, num_layers=6, dropout=0.1):
        super().__init__()
        self.d_model = d_model

        self.audio_input_proj = nn.Linear(d_model, d_model)
        self.pos_embed = SinusoidalPositionalEncoding(d_model, max_len=4096)
        self.dropout = nn.Dropout(dropout)

        # Transformer decoder with causal self-attention + cross-attention to text
        decoder_layer = nn.TransformerDecoderLayer(
            d_model=d_model, nhead=nhead,
            dim_feedforward=d_model * 4,
            dropout=dropout, batch_first=True,
            activation="gelu",
        )
        self.transformer = nn.TransformerDecoder(decoder_layer, num_layers=num_layers)

        # Output projection to predict next embedding
        self.output_proj = nn.Sequential(
            nn.Linear(d_model, d_model),
            nn.GELU(),
            nn.Linear(d_model, d_model),
        )

    def forward(self, audio_emb, text_emb, audio_mask=None, text_mask=None):
        """
        Args:
            audio_emb: [B, T_audio, d_model] — audio embeddings (from encoder)
            text_emb: [B, T_text, d_model] — text embeddings
            audio_mask: [B, T_audio] bool, True = padding
            text_mask: [B, T_text] bool, True = padding
        Returns:
            predicted: [B, T_audio, d_model] — predicted next embeddings
        """
        x = self.audio_input_proj(audio_emb)
        x = self.pos_embed(x)
        x = self.dropout(x)

        # Causal mask: each position can only attend to previous positions
        T = x.shape[1]
        causal_mask = torch.triu(
            torch.ones(T, T, device=x.device, dtype=torch.bool), diagonal=1
        )

        predicted = self.transformer(
            x, text_emb,
            tgt_mask=causal_mask,
            tgt_key_padding_mask=audio_mask,
            memory_key_padding_mask=text_mask,
        )

        predicted = self.output_proj(predicted)
        return predicted

    def _cached_mha(self, mha, q_input, kv_input, cache_k, cache_v):
        """Run multi-head attention with KV cache.
        q_input: [B, 1, d], kv_input: [B, 1, d] (new token for K,V)
        cache_k, cache_v: [B, T_prev, d] projected keys/values
        Returns: output [B, 1, d], updated cache_k, updated cache_v
        """
        d = mha.embed_dim
        nhead = mha.num_heads
        head_dim = d // nhead
        B = q_input.shape[0]

        # Project Q from q_input, K/V from kv_input
        # in_proj_weight is [3d, d] = [Wq; Wk; Wv]
        Wq, Wk, Wv = mha.in_proj_weight.chunk(3, dim=0)
        bq, bk, bv = mha.in_proj_bias.chunk(3, dim=0)

        q = F.linear(q_input, Wq, bq)        # [B, 1, d]
        k_new = F.linear(kv_input, Wk, bk)    # [B, 1, d]
        v_new = F.linear(kv_input, Wv, bv)    # [B, 1, d]

        # Append to cache
        if cache_k is not None and cache_k.shape[1] > 0:
            k = torch.cat([cache_k, k_new], dim=1)
            v = torch.cat([cache_v, v_new], dim=1)
        else:
            k = k_new
            v = v_new

        # Multi-head attention
        q = q.view(B, 1, nhead, head_dim).transpose(1, 2)      # [B, nhead, 1, head_dim]
        k_mh = k.view(B, -1, nhead, head_dim).transpose(1, 2)  # [B, nhead, T, head_dim]
        v_mh = v.view(B, -1, nhead, head_dim).transpose(1, 2)

        attn = torch.matmul(q, k_mh.transpose(-2, -1)) / (head_dim ** 0.5)
        attn = F.softmax(attn, dim=-1)
        out = torch.matmul(attn, v_mh)  # [B, nhead, 1, head_dim]

        out = out.transpose(1, 2).contiguous().view(B, 1, d)
        out = F.linear(out, mha.out_proj.weight, mha.out_proj.bias)

        return out, k, v

    def _cached_cross_attn(self, mha, q_input, memory, cache_k, cache_v, memory_mask=None):
        """Cross-attention with cached memory K,V (computed once).
        memory_mask: [B, T_text] bool, True = padding (applied as -inf)."""
        d = mha.embed_dim
        nhead = mha.num_heads
        head_dim = d // nhead
        B = q_input.shape[0]

        Wq, Wk, Wv = mha.in_proj_weight.chunk(3, dim=0)
        bq, bk, bv = mha.in_proj_bias.chunk(3, dim=0)

        q = F.linear(q_input, Wq, bq)  # [B, 1, d]

        if cache_k is None:
            k = F.linear(memory, Wk, bk)  # [B, T_text, d]
            v = F.linear(memory, Wv, bv)
        else:
            k, v = cache_k, cache_v

        q = q.view(B, 1, nhead, head_dim).transpose(1, 2)
        k_mh = k.view(B, -1, nhead, head_dim).transpose(1, 2)
        v_mh = v.view(B, -1, nhead, head_dim).transpose(1, 2)

        attn = torch.matmul(q, k_mh.transpose(-2, -1)) / (head_dim ** 0.5)
        if memory_mask is not None:
            # memory_mask: [B, T_text] → [B, 1, 1, T_text]
            attn = attn.masked_fill(memory_mask.unsqueeze(1).unsqueeze(2), float('-inf'))
        attn = F.softmax(attn, dim=-1)
        out = torch.matmul(attn, v_mh)

        out = out.transpose(1, 2).contiguous().view(B, 1, d)
        out = F.linear(out, mha.out_proj.weight, mha.out_proj.bias)

        return out, k, v

    def init_cache(self, num_layers):
        """Initialize empty KV cache."""
        return {
            'self_k': [None] * num_layers,
            'self_v': [None] * num_layers,
            'cross_k': [None] * num_layers,
            'cross_v': [None] * num_layers,
        }

    def inference_step(self, new_emb, text_emb, step_idx, cache, text_mask=None):
        """Process one token with KV cache.
        new_emb: [B, 1, d_model] — raw audio embedding (not projected)
        text_mask: [B, T_text] bool, True = padding
        Returns: predicted [B, 1, d_model], updated cache
        """
        x = self.audio_input_proj(new_emb)
        # Add positional encoding for this specific position
        x = x + self.pos_embed.pe[:, step_idx:step_idx + 1]

        for i, layer in enumerate(self.transformer.layers):
            # Self-attention with cache (post-norm)
            sa_out, cache['self_k'][i], cache['self_v'][i] = self._cached_mha(
                layer.self_attn, x, x, cache['self_k'][i], cache['self_v'][i]
            )
            x = layer.norm1(x + layer.dropout1(sa_out))

            # Cross-attention with cached text KV
            ca_out, cache['cross_k'][i], cache['cross_v'][i] = self._cached_cross_attn(
                layer.multihead_attn, x, text_emb, cache['cross_k'][i], cache['cross_v'][i],
                memory_mask=text_mask,
            )
            x = layer.norm2(x + layer.dropout2(ca_out))

            # FFN
            ff_out = layer.linear2(layer.dropout(layer.activation(layer.linear1(x))))
            x = layer.norm3(x + layer.dropout3(ff_out))

        return self.output_proj(x), cache

# ─── Positional Encoding ────────────────────────────────────────────────────

class SinusoidalPositionalEncoding(nn.Module):
    def __init__(self, d_model, max_len=8192):
        super().__init__()
        pe = torch.zeros(max_len, d_model)
        position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)
        div_term = torch.exp(torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model))
        pe[:, 0::2] = torch.sin(position * div_term)
        pe[:, 1::2] = torch.cos(position * div_term)
        self.register_buffer("pe", pe.unsqueeze(0))  # [1, max_len, d_model]

    def forward(self, x):
        return x + self.pe[:, :x.shape[1]]

# ─── Mel Decoder ────────────────────────────────────────────────────────────

class ResConvBlock(nn.Module):
    """Residual conv block for mel decoder."""

    def __init__(self, channels, kernel_size=3):
        super().__init__()
        self.net = nn.Sequential(
            nn.Conv1d(channels, channels, kernel_size, padding=kernel_size // 2),
            nn.GELU(),
            nn.Conv1d(channels, channels, kernel_size, padding=kernel_size // 2),
        )

    def forward(self, x):
        return x + self.net(x)

class MelDecoder(nn.Module):
    """Decodes latent embeddings back to mel spectrograms.
    Improved architecture with residual blocks for sharper mel output."""

    def __init__(self, d_model=256, n_mels=100, upsample_factor=4):
        super().__init__()
        hidden = d_model * 2  # 512 — more capacity

        self.proj = nn.Sequential(
            nn.Linear(d_model, hidden),
            nn.GELU(),
        )
        if upsample_factor == 2:
            self.upsample = nn.Sequential(
                nn.ConvTranspose1d(hidden, hidden, kernel_size=4, stride=2, padding=1),
                nn.GELU(),
                ResConvBlock(hidden, kernel_size=3),
            )
        else:
            self.upsample = nn.Sequential(
                # 4x upsample (2x + 2x)
                nn.ConvTranspose1d(hidden, hidden, kernel_size=4, stride=2, padding=1),
                nn.GELU(),
                ResConvBlock(hidden, kernel_size=3),
                nn.ConvTranspose1d(hidden, hidden, kernel_size=4, stride=2, padding=1),
                nn.GELU(),
                ResConvBlock(hidden, kernel_size=3),
            )
        # Refinement + projection to mel bins
        self.out = nn.Sequential(
            nn.Conv1d(hidden, hidden, kernel_size=7, padding=3),
            nn.GELU(),
            ResConvBlock(hidden, kernel_size=5),
            nn.Conv1d(hidden, n_mels, kernel_size=7, padding=3),
        )

    def forward(self, z):
        """z: [B, T_down, d_model] → mel: [B, n_mels, T_up]"""
        x = self.proj(z)
        x = x.transpose(1, 2)  # [B, hidden, T_down]
        x = self.upsample(x)   # [B, hidden, T_up]
        return self.out(x)      # [B, n_mels, T_up]

# ─── Multi-Resolution Spectral Loss ─────────────────────────────────────────

class MultiResolutionSpectralLoss(nn.Module):
    """Computes spectral convergence + log magnitude loss at multiple STFT resolutions.
    This forces the mel decoder to produce sharp, well-defined spectral structure
    rather than blurry averages that Vocos can't decode properly."""

    def __init__(self, resolutions=((64, 16, 64), (128, 32, 128), (256, 64, 256))):
        super().__init__()
        # Each resolution is (n_fft, hop_length, win_length)
        self.resolutions = resolutions

    def _stft_loss(self, pred, target, n_fft, hop_length, win_length):
        """Compute spectral convergence + log magnitude loss for one resolution."""
        window = torch.hann_window(win_length, device=pred.device)

        # pred/target: [B, n_mels, T] — treat as 1D signal per mel bin
        # Flatten mel bins into batch: [B*n_mels, T]
        B, M, T = pred.shape
        pred_flat = pred.reshape(B * M, T)
        target_flat = target.reshape(B * M, T)

        pred_stft = torch.stft(pred_flat, n_fft, hop_length, win_length,
                               window=window, return_complex=True)
        target_stft = torch.stft(target_flat, n_fft, hop_length, win_length,
                                 window=window, return_complex=True)

        pred_mag = pred_stft.abs()
        target_mag = target_stft.abs()

        # Spectral convergence: Frobenius norm of difference / Frobenius norm of target
        sc_loss = torch.norm(target_mag - pred_mag, p="fro") / (torch.norm(target_mag, p="fro") + 1e-8)

        # Log magnitude loss
        log_mag_loss = F.l1_loss(
            torch.log(pred_mag + 1e-7),
            torch.log(target_mag + 1e-7),
        )

        return sc_loss + log_mag_loss

    def forward(self, pred, target):
        """pred, target: [B, n_mels, T]"""
        loss = torch.tensor(0.0, device=pred.device)
        count = 0
        for n_fft, hop, win in self.resolutions:
            # Skip if T is too short for this resolution
            if pred.shape[2] >= n_fft:
                loss = loss + self._stft_loss(pred, target, n_fft, hop, win)
                count += 1
        if count > 0:
            loss = loss / count
        return loss

# ─── Full LeWM TTS Model ────────────────────────────────────────────────────

class LeWMTTS(nn.Module):
    """
    Complete LeWM TTS model.

    Training: text + mel → audio embeddings → predict next embedding
    Two losses:
        1. Next-embedding prediction: MSE(predicted[t], target[t+1])
        2. Gaussian regularizer: KL(q(z|x) || N(0, I))
    """

    def __init__(self, config):
        super().__init__()
        d_model = config.get("d_model", 256)
        nhead = config.get("nhead", 4)
        n_mels = config.get("n_mels", 100)
        text_vocab = config.get("text_vocab_size", 256)
        text_layers = config.get("text_encoder_layers", 4)
        audio_layers = config.get("audio_encoder_layers", 4)
        predictor_layers = config.get("predictor_layers", 6)
        dropout = config.get("dropout", 0.1)
        self.kl_weight = config.get("kl_weight", 0.01)
        n_speakers = config.get("n_speakers", 1)

        self.text_encoder = TextEncoder(
            vocab_size=text_vocab, d_model=d_model, nhead=nhead,
            num_layers=text_layers, dropout=dropout,
        )
        downsample_factor = config.get("downsample_factor", 4)
        self.audio_encoder = AudioEncoder(
            n_mels=n_mels, d_model=d_model, nhead=nhead,
            num_layers=audio_layers, downsample_factor=downsample_factor,
            dropout=dropout,
        )
        self.predictor = JEPAPredictor(
            d_model=d_model, nhead=nhead,
            num_layers=predictor_layers, dropout=dropout,
        )

        # Mel decoder: reconstructs mel from latent embeddings
        # Needed so the latent space is grounded to actual audio
        self.mel_decoder = MelDecoder(d_model=d_model, n_mels=n_mels, upsample_factor=downsample_factor)
        self.recon_weight = config.get("recon_weight", 1.0)

        # EMA target encoder — core JEPA component
        # Provides slowly-moving prediction targets to prevent representation collapse
        self.ema_audio_encoder = copy.deepcopy(self.audio_encoder)
        for p in self.ema_audio_encoder.parameters():
            p.requires_grad = False
        self.ema_decay = config.get("ema_decay", 0.998)

        # Speaker embedding (for multi-speaker)
        self.n_speakers = n_speakers
        if n_speakers > 1:
            self.speaker_embed = nn.Embedding(n_speakers, d_model)
        else:
            self.speaker_embed = None

        # Learnable start-of-sequence embedding for AR inference
        self.start_emb = nn.Parameter(torch.randn(1, 1, d_model) * 0.02)

        # Prediction weight — primary JEPA loss
        self.pred_weight = config.get("pred_weight", 10.0)

        # Input noise injection to bridge train/inference gap
        self.input_noise = config.get("input_noise", 0.0)

        # Free-bits KL: allow KL up to this threshold without penalty
        self.free_bits = config.get("free_bits", 1.0)

        # Multi-resolution spectral loss for sharper mel output
        self.spectral_loss = MultiResolutionSpectralLoss()
        self.spectral_weight = config.get("spectral_weight", 0.5)

        self.config = config

    def _add_speaker(self, text_emb, speaker_id):
        """Add speaker embedding to text embeddings."""
        if self.speaker_embed is not None and speaker_id is not None:
            spk = self.speaker_embed(speaker_id).unsqueeze(1)  # [B, 1, d]
            text_emb = text_emb + spk
        return text_emb

    @torch.no_grad()
    def update_ema(self):
        """Update EMA target encoder weights."""
        for p_online, p_ema in zip(self.audio_encoder.parameters(),
                                   self.ema_audio_encoder.parameters()):
            p_ema.data.mul_(self.ema_decay).add_(p_online.data, alpha=1 - self.ema_decay)

    def forward(self, mel, text_tokens, mel_mask=None, text_mask=None, speaker_id=None):
        """
        Args:
            mel: [B, n_mels, T_mel]
            text_tokens: [B, T_text]
            mel_mask: [B, T_mel] bool
            text_mask: [B, T_text] bool
        Returns:
            loss_dict with 'prediction_loss', 'kl_loss', 'total_loss'
        """
        # Encode text + speaker conditioning
        text_emb = self.text_encoder(text_tokens, text_mask)
        text_emb = self._add_speaker(text_emb, speaker_id)

        # Online encoder → latent embeddings (used for predictor input + reconstruction)
        z, mu, logvar = self.audio_encoder(mel, mel_mask)
        B = z.shape[0]
        T_down = z.shape[1]

        # EMA target encoder → stable prediction targets (no gradient)
        with torch.no_grad():
            _, mu_ema, _ = self.ema_audio_encoder(mel, mel_mask)
            target_emb = mu_ema  # [B, T_down, d] — all EMA targets

        # Prepend learnable start embedding to input sequence
        # Input:  [start_emb, z[0], z[1], ..., z[T-2]]  → T_down positions
        # Target: [z_ema[0],  z_ema[1], ..., z_ema[T-1]] → T_down positions
        # This trains the model to predict the first audio embedding from start_emb + text
        start = self.start_emb.expand(B, -1, -1)  # [B, 1, d]
        input_emb = torch.cat([start, z[:, :-1]], dim=1)  # [B, T_down, d]

        # Input noise: only on positions 1+ (real embeddings), not start_emb
        # Simulates the imperfect predictions seen during AR inference
        if self.training and self.input_noise > 0:
            noise = torch.zeros_like(input_emb)
            noise[:, 1:] = torch.randn(B, T_down - 1, input_emb.shape[-1],
                                        device=input_emb.device) * self.input_noise
            input_emb = input_emb + noise

        # Adjust masks for downsampled sequence
        if mel_mask is not None:
            ds_mask = mel_mask[:, ::self.audio_encoder.downsample_factor][:, :T_down]
            # Predictor input mask: start_emb (pos 0) is always unmasked
            start_mask = torch.zeros(B, 1, dtype=torch.bool, device=z.device)
            pred_mask = torch.cat([start_mask, ds_mask[:, :-1]], dim=1)  # [B, T_down]
            # Loss mask: follows original ds_mask for target positions
            loss_mask = ds_mask
        else:
            pred_mask = None
            loss_mask = None

        predicted = self.predictor(input_emb, text_emb, pred_mask, text_mask)

        # ─── Loss 1: Next-embedding prediction (MSE) against EMA targets ───
        if loss_mask is not None:
            valid = (~loss_mask).unsqueeze(-1)  # [B, T_down, 1]
            prediction_loss = F.mse_loss(
                predicted * valid, target_emb * valid, reduction="sum"
            ) / (valid.sum() * predicted.shape[-1] + 1e-8)
        else:
            prediction_loss = F.mse_loss(predicted, target_emb)

        # ─── Loss 2: Free-bits KL — only penalize KL above threshold ───
        # Prevents both KL collapse (too low) and KL explosion (too high)
        kl_per_dim = -0.5 * (1 + logvar - mu.pow(2) - logvar.exp())  # [B, T, d]
        kl_free = torch.clamp(kl_per_dim - self.free_bits, min=0.0)
        if mel_mask is not None:
            T_down = mu.shape[1]
            ds_mask = mel_mask[:, ::self.audio_encoder.downsample_factor][:, :T_down]
            valid_kl = (~ds_mask).unsqueeze(-1).float()
            kl_loss = (kl_free * valid_kl).sum() / (valid_kl.sum() * mu.shape[-1] + 1e-8)
        else:
            kl_loss = kl_free.mean()

        # ─── Loss 3: Mel reconstruction ───
        mel_recon = self.mel_decoder(z)  # [B, n_mels, T_up]
        T_mel = mel.shape[2]
        T_recon = mel_recon.shape[2]
        if T_recon > T_mel:
            mel_recon = mel_recon[:, :, :T_mel]
        elif T_recon < T_mel:
            mel = mel[:, :, :T_recon]
            if mel_mask is not None:
                mel_mask = mel_mask[:, :T_recon]

        if mel_mask is not None:
            valid_mel = (~mel_mask).unsqueeze(1).float()  # [B, 1, T]
            recon_loss = F.l1_loss(
                mel_recon * valid_mel, mel * valid_mel, reduction="sum"
            ) / (valid_mel.sum() * mel.shape[1] + 1e-8)
        else:
            recon_loss = F.l1_loss(mel_recon, mel)

        # ─── Loss 4: Multi-resolution spectral loss ───
        if mel_mask is not None:
            spectral_loss = self.spectral_loss(mel_recon * valid_mel, mel * valid_mel)
        else:
            spectral_loss = self.spectral_loss(mel_recon, mel)

        total_loss = (self.pred_weight * prediction_loss
                      + self.kl_weight * kl_loss
                      + self.recon_weight * recon_loss
                      + self.spectral_weight * spectral_loss)

        return {
            "total_loss": total_loss,
            "prediction_loss": prediction_loss,
            "kl_loss": kl_loss,
            "recon_loss": recon_loss,
            "spectral_loss": spectral_loss,
        }

    def encode_audio(self, mel):
        """Encode mel to latent embeddings (for inference)."""
        z, mu, _ = self.audio_encoder(mel)
        return mu  # Use mean at inference

    def predict_next(self, audio_emb, text_emb, text_mask=None):
        """Predict next audio embedding autoregressively."""
        predicted = self.predictor(audio_emb, text_emb, text_mask=text_mask)
        return predicted[:, -1:]  # Only the last prediction

    def init_ar_cache(self):
        """Initialize KV cache for autoregressive inference."""
        num_layers = len(self.predictor.transformer.layers)
        return self.predictor.init_cache(num_layers)

    def predict_next_cached(self, new_emb, text_emb, step_idx, cache, text_mask=None):
        """Predict next embedding using KV cache. O(1) per step instead of O(n)."""
        return self.predictor.inference_step(new_emb, text_emb, step_idx, cache, text_mask=text_mask)

def count_parameters(model):
    return sum(p.numel() for p in model.parameters() if p.requires_grad)

def build_model(config=None):
    if config is None:
        config = {
            "d_model": 256,
            "nhead": 4,
            "n_mels": 100,
            "text_vocab_size": 256,
            "text_encoder_layers": 4,
            "audio_encoder_layers": 4,
            "predictor_layers": 6,
            "dropout": 0.1,
            "kl_weight": 0.05,
            "pred_weight": 10.0,
            "ema_decay": 0.998,
            "free_bits": 1.0,
            "input_noise": 0.0,
        }
    model = LeWMTTS(config)
    print(f"LeWM TTS model: {count_parameters(model)/1e6:.2f}M parameters")
    return model, config

+        # Force eval mode so dropout is disabled — EMA targets must be deterministic

+            self.ema_audio_encoder.eval()

+            self.ema_audio_encoder.train()


Now fix inference_engine.py — replace random start embeddings with learned ones:


        # Precompute start embedding scale
        self._start_scale = 5.0

        print(f"Engine ready: {checkpoint_path}")


The file /home/ubuntu/lewm-tts/inference_engine.py has been updated successfully.

"""
Production inference engine for LeWM TTS.
Features:
  - Batched inference (many concurrent requests in parallel)
  - KV-cache (O(1) per step instead of O(n))
  - FP16
  - Audio prompt support (for multi-speaker / voice cloning)
  - Continuous batching (add/remove requests dynamically)
  - Griffin-Lim + Vocos vocoder pipeline
"""

import torch
import torch.nn.functional as F
import torchaudio
import numpy as np
import time
from dataclasses import dataclass, field
from typing import Optional
from pathlib import Path

from model import LeWMTTS

# ─── Vocoder ────────────────────────────────────────────────────────────────

class Vocoder:
    """Vocos vocoder for 100-mel → waveform (direct decode, no Griffin-Lim)."""

    def __init__(self):
        from vocos import Vocos
        self.vocos = Vocos.from_pretrained("charactr/vocos-mel-24khz")
        self.vocos.eval()

    def __call__(self, mel):
        """mel: [B, 100, T] → waveforms list of [T_audio] numpy arrays."""
        with torch.no_grad():
            # Vocos decode handles batched input
            waveforms = []
            for i in range(mel.shape[0]):
                wav = self.vocos.decode(mel[i:i+1])  # [1, T_audio]
                waveforms.append(wav.squeeze(0).numpy())
        return waveforms

# ─── Request / Result ───────────────────────────────────────────────────────

@dataclass
class TTSRequest:
    text: str
    max_steps: int = 300
    temperature: float = 0.0
    prompt_mel: Optional[torch.Tensor] = None  # [1, 80, T] for voice cloning

@dataclass
class TTSResult:
    waveform: np.ndarray
    sample_rate: int = 24000
    steps_taken: int = 0
    generation_time: float = 0.0

    def save(self, path):
        wav = torch.from_numpy(self.waveform).unsqueeze(0)
        torchaudio.save(path, wav, self.sample_rate)

    @property
    def duration(self):
        return len(self.waveform) / self.sample_rate

# ─── Inference Engine ───────────────────────────────────────────────────────

class InferenceEngine:
    """
    High-throughput batched TTS inference engine.

    Usage:
        engine = InferenceEngine("checkpoint.pt")

        # Single request
        result = engine.synthesize("नमस्ते भारत")

        # Batched (high throughput)
        results = engine.synthesize_batch(["text1", "text2", ...])

        # With voice prompt
        result = engine.synthesize("text", prompt_audio="ref.wav")

        # Benchmark
        engine.benchmark()
    """

    def __init__(self, checkpoint_path, device="cuda", dtype=torch.float16):
        self.device = torch.device(device if torch.cuda.is_available() else "cpu")
        self.dtype = dtype

        # Load model
        ckpt = torch.load(checkpoint_path, map_location=self.device, weights_only=False)
        self.config = ckpt["config"]
        self.model = LeWMTTS(self.config).to(self.device).to(dtype)
        self.model.load_state_dict(ckpt["model"])
        self.model.eval()

        self.d_model = self.config["d_model"]

        # Vocoder (runs on CPU)
        self.vocoder = Vocoder()

        # Precompute start embedding scale
        self._start_scale = 5.0

        print(f"Engine ready: {checkpoint_path}")
        print(f"  Device: {self.device}, Dtype: {dtype}")
        print(f"  Model: {sum(p.numel() for p in self.model.parameters())/1e6:.1f}M params")

    # ─── Text encoding ──────────────────────────────────────────────────

    def _encode_texts(self, texts):
        """Encode list of texts → padded embeddings + mask.
        Returns: text_emb [B, T_max, d], text_mask [B, T_max] (True=padding)
        """
        # Byte-level tokenization
        token_lists = [list(t.encode("utf-8")) for t in texts]
        max_len = max(len(t) for t in token_lists)

        # Pad to same length
        padded = torch.zeros(len(texts), max_len, dtype=torch.long, device=self.device)
        mask = torch.ones(len(texts), max_len, dtype=torch.bool, device=self.device)

        for i, tokens in enumerate(token_lists):
            padded[i, :len(tokens)] = torch.tensor(tokens, dtype=torch.long)
            mask[i, :len(tokens)] = False

        with torch.no_grad():
            text_emb = self.model.text_encoder(padded, text_mask=mask)

        return text_emb, mask

    # ─── Audio prompt encoding ──────────────────────────────────────────

    def _encode_prompt(self, prompt_mel):
        """Encode prompt mel → embeddings for AR seeding.
        prompt_mel: [1, 80, T]
        Returns: prompt_embs [1, T_down, d]
        """
        prompt_mel = prompt_mel.to(self.device).to(self.dtype)
        with torch.no_grad():
            embs = self.model.encode_audio(prompt_mel)  # [1, T_down, d]
        return embs

    def load_prompt_audio(self, audio_path, max_seconds=3.0):
        """Load audio file → mel for prompting.
        Returns: mel [1, 80, T]
        """
        wav, sr = torchaudio.load(audio_path)
        if sr != 24000:
            wav = torchaudio.functional.resample(wav, sr, 24000)
        # Trim to max_seconds
        max_samples = int(max_seconds * 24000)
        wav = wav[:1, :max_samples]  # mono, trimmed

        mel_transform = torchaudio.transforms.MelSpectrogram(
            sample_rate=24000, n_fft=1024, hop_length=256,
            n_mels=100, power=1.0,
        )
        mel = mel_transform(wav)
        mel = torch.log(mel.clamp(min=1e-7))
        return mel  # [1, 80, T]

    # ─── Core AR generation (batched) ───────────────────────────────────

    def _generate_batch(self, text_emb, text_mask, batch_size,
                        max_steps=300, temperature=0.0,
                        prompt_embs=None):
        """
        Core batched AR generation with KV-cache.

        Args:
            text_emb: [B, T_text, d]
            text_mask: [B, T_text]
            batch_size: int
            max_steps: int
            temperature: float
            prompt_embs: list of [1, T_prompt, d] or None per request

        Returns:
            all_embs: [B, T_total, d]
            steps_taken: list of int per request
        """
        d = self.d_model
        cache = self.model.init_ar_cache()

        # Track per-request state
        active = torch.ones(batch_size, dtype=torch.bool, device=self.device)
        steps_taken = [0] * batch_size
        all_emb_lists = [[] for _ in range(batch_size)]
        prev_norms = [[] for _ in range(batch_size)]

        # Handle prompts: seed cache with prompt embeddings
        if prompt_embs is not None:
            max_prompt_len = max(
                (e.shape[1] for e in prompt_embs if e is not None), default=0
            )
            if max_prompt_len > 0:
                # Process prompt embeddings through cache step by step
                # Pad prompts to same length
                prompt_padded = torch.zeros(
                    batch_size, max_prompt_len, d,
                    device=self.device, dtype=self.dtype,
                )
                prompt_lengths = []
                for i, pe in enumerate(prompt_embs):
                    if pe is not None:
                        L = pe.shape[1]
                        prompt_padded[i, :L] = pe[0]
                        prompt_lengths.append(L)
                    else:
                        prompt_lengths.append(0)

                # Feed prompt through cache one step at a time
                for t in range(max_prompt_len):
                    step_emb = prompt_padded[:, t:t+1, :]  # [B, 1, d]
                    _, cache = self.model.predict_next_cached(
                        step_emb, text_emb, t, cache, text_mask=text_mask,
                    )
                    # Store embeddings
                    for i in range(batch_size):
                        if t < prompt_lengths[i]:
                            all_emb_lists[i].append(step_emb[i:i+1])

                step_offset = max_prompt_len
                # Use last prompt prediction as first AR input
                last_pred, cache = self.model.predict_next_cached(
                    prompt_padded[:, -1:, :], text_emb, step_offset, cache,
                    text_mask=text_mask,
                )
                current_emb = last_pred
                step_offset += 1
            else:
                step_offset = 0
                current_emb = None
        else:
            step_offset = 0
            current_emb = None

        # If no prompt, start with random embedding
        if current_emb is None:
            start_emb = torch.randn(
                batch_size, 1, d, device=self.device, dtype=self.dtype,
            ) * self._start_scale
            current_emb, cache = self.model.predict_next_cached(
                start_emb, text_emb, 0, cache, text_mask=text_mask,
            )
            for i in range(batch_size):
                all_emb_lists[i].append(start_emb[i:i+1])
            step_offset = 1

        if temperature > 0:
            current_emb = current_emb + torch.randn_like(current_emb) * temperature

        for i in range(batch_size):
            all_emb_lists[i].append(current_emb[i:i+1])

        # AR loop
        for step in range(1, max_steps):
            if not active.any():
                break

            next_emb, cache = self.model.predict_next_cached(
                current_emb, text_emb, step + step_offset, cache,
                text_mask=text_mask,
            )

            if temperature > 0:
                next_emb = next_emb + torch.randn_like(next_emb) * temperature

            # Store and check stopping per request
            for i in range(batch_size):
                if not active[i]:
                    continue

                all_emb_lists[i].append(next_emb[i:i+1])
                steps_taken[i] = step

                norm = next_emb[i].norm().item()
                prev_norms[i].append(norm)
                if len(prev_norms[i]) > 20:
                    recent = prev_norms[i][-20:]
                    if np.std(recent) < 0.01 * np.mean(recent) and step > 30:
                        active[i] = False

            current_emb = next_emb

        # Concatenate embeddings per request
        results = []
        for i in range(batch_size):
            embs = torch.cat(all_emb_lists[i], dim=1)  # [1, T, d]
            results.append(embs)

        return results, steps_taken

    # ─── Public API ─────────────────────────────────────────────────────

    def synthesize(self, text, max_steps=300, temperature=0.0,
                   prompt_audio=None, prompt_seconds=3.0):
        """Synthesize single text.

        Args:
            text: Hindi text string
            max_steps: max AR steps
            temperature: sampling temperature (0 = greedy)
            prompt_audio: path to reference audio for voice cloning
            prompt_seconds: max seconds of prompt to use

        Returns: TTSResult
        """
        return self.synthesize_batch(
            [text], max_steps=max_steps, temperature=temperature,
            prompt_audios=[prompt_audio] if prompt_audio else None,
            prompt_seconds=prompt_seconds,
        )[0]

    def synthesize_batch(self, texts, max_steps=300, temperature=0.0,
                         prompt_audios=None, prompt_seconds=3.0):
        """Synthesize a batch of texts in parallel.

        Args:
            texts: list of text strings
            max_steps: max AR steps
            temperature: sampling temperature
            prompt_audios: list of audio paths (or None) per text
            prompt_seconds: max seconds of prompt

        Returns: list of TTSResult
        """
        B = len(texts)
        t0 = time.time()

        # Encode texts
        text_emb, text_mask = self._encode_texts(texts)

        # Encode prompts if given
        prompt_embs = None
        if prompt_audios is not None:
            prompt_embs = []
            for pa in prompt_audios:
                if pa is not None:
                    mel = self.load_prompt_audio(pa, max_seconds=prompt_seconds)
                    embs = self._encode_prompt(mel)
                    prompt_embs.append(embs)
                else:
                    prompt_embs.append(None)

        # Generate
        with torch.no_grad():
            emb_results, steps_taken = self._generate_batch(
                text_emb, text_mask, B,
                max_steps=max_steps, temperature=temperature,
                prompt_embs=prompt_embs,
            )

        # Decode to mel → waveform
        results = []
        for i in range(B):
            with torch.no_grad():
                mel = self.model.mel_decoder(emb_results[i].to(self.device))
                # mel: [1, 80, T]
            wav_list = self.vocoder(mel.float().cpu())
            waveform = wav_list[0]

            # Normalize
            mx = np.abs(waveform).max()
            if mx > 0:
                waveform = waveform / mx * 0.95

            results.append(TTSResult(
                waveform=waveform,
                steps_taken=steps_taken[i],
                generation_time=time.time() - t0,
            ))

        return results

    # ─── Benchmark ──────────────────────────────────────────────────────

    def benchmark(self, batch_sizes=None, steps=300):
        """Benchmark throughput at various batch sizes."""
        if batch_sizes is None:
            batch_sizes = [1, 4, 16, 32, 64, 128, 256]

        text = "भारत एक महान देश है और हम सब भारतवासी हैं।"
        audio_per_step = 4 * 256 / 24000  # seconds of audio per AR step

        print(f"\n{'='*70}")
        print(f"  LeWM TTS Inference Benchmark — {self.device}, {self.dtype}")
        print(f"  Model: {self.config['d_model']}d, {self.config['predictor_layers']}L predictor")
        print(f"  AR steps: {steps}")
        print(f"{'='*70}")
        print(f"  {'Batch':>6} {'Time(s)':>8} {'Audio(s)':>9} {'RTF':>8} "
              f"{'Speed':>8} {'Throughput':>12} {'Per-req':>10}")
        print(f"  {'-'*6} {'-'*8} {'-'*9} {'-'*8} {'-'*8} {'-'*12} {'-'*10}")

        for B in batch_sizes:
            try:
                texts = [text] * B
                token_lists = [list(t.encode("utf-8")) for t in texts]
                max_len = max(len(t) for t in token_lists)
                padded = torch.zeros(B, max_len, dtype=torch.long, device=self.device)
                mask = torch.ones(B, max_len, dtype=torch.bool, device=self.device)
                for i, tokens in enumerate(token_lists):
                    padded[i, :len(tokens)] = torch.tensor(tokens, dtype=torch.long)
                    mask[i, :len(tokens)] = False

                with torch.no_grad():
                    text_emb = self.model.text_encoder(padded, text_mask=mask)

                # Warmup
                if B == batch_sizes[0]:
                    cache = self.model.init_ar_cache()
                    start = torch.randn(B, 1, self.d_model,
                                        device=self.device, dtype=self.dtype) * self._start_scale
                    nxt = start
                    for s in range(5):
                        nxt, cache = self.model.predict_next_cached(
                            nxt, text_emb, s, cache, text_mask=mask)
                    torch.cuda.synchronize()

                # Benchmark
                cache = self.model.init_ar_cache()
                start = torch.randn(B, 1, self.d_model,
                                    device=self.device, dtype=self.dtype) * self._start_scale
                nxt = start

                torch.cuda.synchronize()
                t0 = time.time()
                for s in range(steps):
                    nxt, cache = self.model.predict_next_cached(
                        nxt, text_emb, s, cache, text_mask=mask)
                torch.cuda.synchronize()
                elapsed = time.time() - t0

                total_audio = B * steps * audio_per_step
                rtf = elapsed / (steps * audio_per_step)
                speed = (steps * audio_per_step) / elapsed
                throughput = total_audio / elapsed
                per_req = elapsed / B

                print(f"  {B:>6} {elapsed:>8.3f} {total_audio:>9.1f} {rtf:>8.4f} "
                      f"{speed:>7.0f}x {throughput:>10.0f}x {per_req:>9.3f}s")

            except torch.cuda.OutOfMemoryError:
                print(f"  {B:>6} {'OOM':>8}")
                torch.cuda.empty_cache()
                break

        print(f"{'='*70}")
        print(f"  Throughput = total audio-seconds generated per wall-second")
        print(f"  Speed = single-request real-time factor")
        print(f"  Per-req = wall-time per request in batch")
        print()

# ─── CLI ────────────────────────────────────────────────────────────────────

if __name__ == "__main__":
    import argparse

    parser = argparse.ArgumentParser(description="LeWM TTS Inference Engine")
    parser.add_argument("--checkpoint", required=True)
    parser.add_argument("--text", default=None, help="Text to synthesize")
    parser.add_argument("--texts", nargs="+", help="Multiple texts for batch synthesis")
    parser.add_argument("--prompt", default=None, help="Reference audio for voice cloning")
    parser.add_argument("--output", default="output.wav")
    parser.add_argument("--output_dir", default=None, help="Output dir for batch mode")
    parser.add_argument("--max_steps", type=int, default=300)
    parser.add_argument("--temperature", type=float, default=0.0)
    parser.add_argument("--benchmark", action="store_true", help="Run throughput benchmark")
    parser.add_argument("--device", default="cuda")
    parser.add_argument("--fp32", action="store_true", help="Use FP32 instead of FP16")
    args = parser.parse_args()

    dtype = torch.float32 if args.fp32 else torch.float16
    engine = InferenceEngine(args.checkpoint, device=args.device, dtype=dtype)

    if args.benchmark:
        engine.benchmark()

    elif args.texts:
        results = engine.synthesize_batch(
            args.texts, max_steps=args.max_steps, temperature=args.temperature,
            prompt_audios=[args.prompt] * len(args.texts) if args.prompt else None,
        )
        out_dir = Path(args.output_dir or "output")
        out_dir.mkdir(exist_ok=True)
        for i, r in enumerate(results):
            path = out_dir / f"batch_{i:03d}.wav"
            r.save(str(path))
            print(f"  [{i}] {path} — {r.duration:.2f}s, {r.steps_taken} steps")

    elif args.text:
        result = engine.synthesize(
            args.text, max_steps=args.max_steps, temperature=args.temperature,
            prompt_audio=args.prompt,
        )
        result.save(args.output)
        print(f"Saved: {args.output} — {result.duration:.2f}s, "
              f"{result.steps_taken} steps, {result.generation_time:.2f}s")
    else:
        parser.print_help()

-        # Precompute start embedding scale


        # If no prompt, start with random embedding
        if current_emb is None:
            start_emb = torch.randn(
                batch_size, 1, d, device=self.device, dtype=self.dtype,
            ) * self._start_scale
            current_emb, cache = self.model.predict_next_cached(
                start_emb, text_emb, 0, cache, text_mask=text_mask,
            )
            for i in range(batch_size):
                all_emb_lists[i].append(start_emb[i:i+1])
            step_offset = 1

        # If no prompt, start with learned start embedding
        if current_emb is None:
            start_emb = self.model.start_emb.expand(batch_size, -1, -1).to(self.dtype)
            current_emb, cache = self.model.predict_next_cached(
                start_emb, text_emb, 0, cache, text_mask=text_mask,
            )
            for i in range(batch_size):
                all_emb_lists[i].append(start_emb[i:i+1])
            step_offset = 1