""" 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, )