verify_model.py

"""
ViuAI_TTS_200M Model Verification Script.
Checks architecture parameter budget, Human Prosody Engine, CFG support, and Voice Prompting.
"""

import sys
import os
import torch

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

# Ensure local package path is recognized
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))

from models import ViuAITTS200M


def run_verification():
    print("=" * 65)
    print("[*] Initializing ViuAI_TTS_200M Architecture...")
    print("=" * 65)

    # Instantiate model
    model = ViuAITTS200M()
    model.eval()

    # 1. Check parameter counts
    param_counts = model.count_parameters()
    print("\n[+] Model Parameter Budget Breakdown:")
    print("-" * 55)
    for name, count in param_counts.items():
        million_params = count / 1e6
        print(f"  {name:.<40} : {million_params:>6.2f} Million ({count:,} params)")
    print("-" * 55)

    # 2. Test Training Step (Flow Matching + Human Prosody Losses + Masking)
    print("\n[+] Testing Training Step (Flow Matching + Human Prosody Losses + Masking)...")
    batch_size = 2
    text_len = 15
    mel_len = 80

    dummy_tokens = torch.randint(0, 200, (batch_size, text_len))
    dummy_mels = torch.randn(batch_size, 80, mel_len)
    dummy_durations = torch.randint(2, 8, (batch_size, text_len))
    dummy_pitch = torch.randn(batch_size, text_len)
    dummy_energy = torch.randn(batch_size, text_len)
    dummy_text_lens = torch.tensor([15, 10], dtype=torch.long)
    dummy_mel_lens = torch.tensor([80, 60], dtype=torch.long)

    dummy_ref_mel = torch.randn(batch_size, 80, 40)

    losses = model.forward_train(
        text_tokens=dummy_tokens,
        mel_target=dummy_mels,
        durations=dummy_durations,
        pitch_target=dummy_pitch,
        energy_target=dummy_energy,
        ref_mel=dummy_ref_mel,
        text_lengths=dummy_text_lens,
        mel_lengths=dummy_mel_lens,
        train_vocoder=False,
    )

    print(f"  [OK] Flow Matching Loss : {losses['flow_loss'].item():.4f}")
    print(f"  [OK] Prosody Loss       : {losses['prosody_loss'].item():.4f}")
    print(f"  [OK] Duration Loss      : {losses['duration_loss'].item():.4f}")
    print(f"  [OK] Total Loss         : {losses['total_loss'].item():.4f}")

    # Gradient verification on DiT and Text Encoder
    losses["total_loss"].backward()
    dit_grads = sum(1 for p in model.dit.parameters() if p.grad is not None and p.grad.abs().sum() > 0)
    total_dit = len(list(model.dit.parameters()))
    print(f"  [OK] DiT Gradients      : {dit_grads}/{total_dit} parameters actively updated")
    model.zero_grad()

    # 3. Test Inference Pass with CFG and Reference Voice Prompting
    print("\n[+] Testing Human-like Inference Pass (with CFG=2.0 & Voice Prompting)...")
    single_token_input = torch.randint(0, 200, (1, 10))
    dummy_ref_audio = torch.randn(1, 80, 50)  # 50 frames of reference audio (like Adam)

    with torch.no_grad():
        audio_out, mel_out = model.generate(
            text_tokens=single_token_input,
            ref_mel=dummy_ref_audio,
            num_steps=10,
            cfg_scale=2.0,
            speed=1.0,
        )

    print(f"  [OK] Generated Mel-Spectrogram Shape : {list(mel_out.shape)}")
    print(f"  [OK] Generated 24kHz Audio Wave Shape : {list(audio_out.shape)}")
    print(f"  [OK] Total Audio Samples Generated   : {audio_out.shape[-1]:,} samples")
    print(f"  [OK] Audio Duration                  : {audio_out.shape[-1] / 24000:.2f} seconds")

    # 4. Test Checkpoint Strict Loading
    ckpt_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), "checkpoints", "viuai_tts_200m_init.pt")
    if os.path.exists(ckpt_path):
        print("\n[+] Testing Checkpoint Compatibility (strict=True)...")
        try:
            ckpt = torch.load(ckpt_path, map_location="cpu", weights_only=True)
        except Exception:
            ckpt = torch.load(ckpt_path, map_location="cpu", weights_only=False)
        state_dict = ckpt["model_state_dict"] if "model_state_dict" in ckpt else ckpt
        model.load_state_dict(state_dict, strict=True)
        print("  [OK] Checkpoint passed strict=True loading verification!")

    print("\n" + "=" * 65)
    print("[SUCCESS] All ViuAI_TTS_200M Human-Realism checks passed successfully!")
    print("=" * 65)


if __name__ == "__main__":
    run_verification()