inference.py

"""
Inference & Audio Generation Script for ViuAI_TTS_200M.
Enhanced with:
  1. In-Context Human Voice Cloning (--ref_audio sample.wav)
  2. Classifier-Free Guidance (--cfg_scale 2.0)
  3. High-Fidelity 15-step Flow Matching ODE solver
"""

import sys
import os
import wave
import torch
import argparse
from typing import Optional

# Setup encoding on Windows
if sys.platform == "win32":
    try:
        sys.stdout.reconfigure(encoding="utf-8")
    except Exception:
        pass

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))

from models import ViuAITTS200M
from dataset.tokenizer import ViuAITokenizer
from dataset.audio_processing import extract_mel_from_file


def load_wav_as_mel(wav_path: str, target_sr: int = 24000) -> Optional[torch.Tensor]:
    """
    Loads reference voice, resamples to 24kHz, and extracts 80-channel log-mel spectrogram.
    """
    if not os.path.exists(wav_path):
        print(f"[!] Warning: Reference audio not found at {wav_path}. Proceeding without reference prompt.")
        return None
    return extract_mel_from_file(wav_path, target_sr=target_sr)


def synthesize_speech(
    text: str,
    output_wav_path: str = "output_speech.wav",
    ref_audio_path: Optional[str] = None,
    checkpoint_path: Optional[str] = None,
    num_steps: int = 15,
    cfg_scale: float = 2.0,
    speed: float = 1.0,
    seed: Optional[int] = 42,
    allow_random_weights: bool = False,
    device: Optional[str] = None,
):
    if seed is not None:
        torch.manual_seed(seed)
        if torch.cuda.is_available():
            torch.cuda.manual_seed_all(seed)

    if device is None or device == "auto":
        dev = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    else:
        dev = torch.device(device)
    device = dev
    print(f"[*] Synthesizing text on device: {device}")
    print(f"[*] Input Text: '{text}'")
    print(f"[*] Speaking Rate (Speed): {speed}x")
    print(f"[*] CFG Guidance Scale: {cfg_scale}x (Studio Broadcast Presence)")
    print(f"[*] Flow Matching Steps: {num_steps} ODE steps")

    # 1. Initialize tokenizer and encode text
    tokenizer = ViuAITokenizer()
    token_ids = tokenizer.text_to_ids(text)
    token_tensor = torch.tensor([token_ids], dtype=torch.long, device=device)

    # 2. Reference audio prompt (Voice Cloning like Adam)
    ref_mel = None
    if ref_audio_path:
        print(f"[*] Loading human reference voice prompt: {ref_audio_path}")
        ref_mel = load_wav_as_mel(ref_audio_path, target_sr=24000)
        if ref_mel is not None:
            ref_mel = ref_mel.to(device)
            print(f"  [OK] Extracted Reference Mel Features: {list(ref_mel.shape)}")

    # 3. Initialize model
    model = ViuAITTS200M().to(device)
    if checkpoint_path:
        if not os.path.exists(checkpoint_path):
            raise FileNotFoundError(f"[!] Checkpoint file not found: {checkpoint_path}")

        print(f"[*] Loading checkpoint: {checkpoint_path}")
        try:
            ckpt = torch.load(checkpoint_path, map_location=device, weights_only=True)
        except Exception:
            ckpt = torch.load(checkpoint_path, map_location=device, weights_only=False)
        state_dict = ckpt.get("ema_state_dict", ckpt.get("model_state_dict", ckpt))
        if "ema_state_dict" in ckpt:
            print("  [OK] Loaded EMA smoothed weights for high-fidelity synthesis!")

        # Auto-remap legacy duration_predictor keys if present
        remapped = {}
        for k, v in state_dict.items():
            if k.startswith("duration_predictor.blocks."):
                remapped[k.replace("duration_predictor.blocks.", "prosody_engine.shared_blocks.")] = v
            elif k.startswith("duration_predictor.proj."):
                remapped[k.replace("duration_predictor.proj.", "prosody_engine.duration_head.")] = v
            elif k.startswith("duration_predictor.head."):
                remapped[k.replace("duration_predictor.head.", "prosody_engine.duration_head.")] = v
            else:
                remapped[k] = v

        try:
            model.load_state_dict(remapped, strict=True)
            print("  [OK] Checkpoint verified and loaded successfully (strict=True)!")
        except RuntimeError as e:
            print(f"  [!] Strict loading warning: {e}")
            missing, unexpected = model.load_state_dict(remapped, strict=False)
            if missing:
                print(f"  [!] Missing keys ({len(missing)}): {missing[:5]}")
            if unexpected:
                print(f"  [!] Unexpected keys ({len(unexpected)}): {unexpected[:5]}")
    elif not allow_random_weights:
        raise ValueError(
            "No checkpoint provided! ViuAI_TTS requires a trained checkpoint to synthesize audio. "
            "Specify --checkpoint path/to/model.pt or pass --allow_random_weights for architecture test."
        )
    else:
        print("[!] Running with randomly initialized weights (--allow_random_weights enabled).")

    model.eval()

    # 4. Generate audio using Flow Matching with CFG
    print(f"[*] Generating speech via Flow Matching ODE ({num_steps} steps, CFG={cfg_scale})...")
    with torch.no_grad():
        audio_tensor, mel_spec = model.generate(
            text_tokens=token_tensor,
            ref_mel=ref_mel,
            num_steps=num_steps,
            cfg_scale=cfg_scale,
            speed=speed,
        )

    # 5. Save output to standard 24kHz 16-bit PCM .wav with DC Removal + Peak Normalization
    audio_flat = audio_tensor.squeeze().detach().cpu().float()

    # Step A: DC Offset Removal (Removes massive electrical hum/bias)
    audio_flat = audio_flat - audio_flat.mean()

    # Step B: High-pass filtering (Remove sub-audible rumble below 50Hz)
    try:
        import torchaudio
        audio_flat = torchaudio.functional.highpass_biquad(audio_flat.unsqueeze(0), 24000, cutoff_freq=50.0).squeeze(0)
        audio_flat = audio_flat - audio_flat.mean()
    except Exception:
        pass

    # Step C: Safe peak normalization to -1 dB (0.89) without amplifying DC bias
    max_amp = audio_flat.abs().max().item()
    if max_amp > 1e-5:
        audio_flat = audio_flat * (0.89 / max_amp)
    audio_flat = audio_flat.clamp(-1.0, 1.0)
    audio_int16 = (audio_flat * 32767.0).to(torch.int16).numpy().tobytes()

    with wave.open(output_wav_path, "wb") as wav_out:
        wav_out.setnchannels(1)        # Mono
        wav_out.setsampwidth(2)        # 16-bit
        wav_out.setframerate(24000)    # 24,000 Hz
        wav_out.writeframes(audio_int16)

    duration_sec = len(audio_flat) / 24000.0
    print(f"[SUCCESS] Audio saved to: {output_wav_path}")
    print(f"  - Sample Rate: 24,000 Hz")
    print(f"  - Audio Length: {duration_sec:.2f} seconds")
    print(f"  - Mel Shape: {list(mel_spec.shape)}")
    return output_wav_path


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="ViuAI_TTS_200M Human Speech Synthesis")
    parser.add_argument(
        "--text",
        type=str,
        default="Namaste! Hum is model ko bilkul natural human aawaz ki tarah bolne ke liye tayar kar rahe hain.",
        help="Text to synthesize",
    )
    parser.add_argument(
        "--ref_audio",
        type=str,
        default=None,
        help="Path to 3-5 sec human reference audio .wav (e.g. Adam voice sample) for zero-shot cloning",
    )
    parser.add_argument(
        "--checkpoint",
        type=str,
        default="checkpoints/viuai_tts_200m_init.pt",
        help="Path to trained checkpoint (.pt)",
    )
    parser.add_argument(
        "--output",
        type=str,
        default="human_voice_test.wav",
        help="Output .wav filename",
    )
    parser.add_argument(
        "--steps",
        type=int,
        default=15,
        help="Flow Matching ODE steps (15-25 for ultra-human realism)",
    )
    parser.add_argument(
        "--cfg_scale",
        type=float,
        default=2.0,
        help="Classifier-Free Guidance scale (1.5 - 2.5)",
    )
    parser.add_argument(
        "--speed",
        type=float,
        default=1.0,
        help="Speaking speed multiplier (e.g. 1.0 = normal, 1.2 = faster, 0.8 = slower)",
    )
    parser.add_argument(
        "--seed",
        type=int,
        default=42,
        help="Random seed for reproducible synthesis",
    )
    parser.add_argument(
        "--allow_random_weights",
        action="store_true",
        help="Allow running without a valid checkpoint (outputs test noise)",
    )
    parser.add_argument(
        "--device",
        type=str,
        default="auto",
        help="Device to run inference on: 'cuda', 'cpu', or 'auto' (default)",
    )
    args = parser.parse_args()

    synthesize_speech(
        text=args.text,
        output_wav_path=args.output,
        ref_audio_path=args.ref_audio,
        checkpoint_path=args.checkpoint if (args.checkpoint and os.path.exists(args.checkpoint)) else (args.checkpoint if not args.allow_random_weights else None),
        num_steps=args.steps,
        cfg_scale=args.cfg_scale,
        speed=args.speed,
        seed=args.seed,
        allow_random_weights=args.allow_random_weights,
        device=args.device,
    )