test_modules/test_latency_optimizations.py

"""
Test latency optimizations in voice_clone.py and tts.py.

Tests configuration logic, dtype selection, attention selection,
audio trimming, and generation parameter values — without loading
full models (no GPU required).
"""
import os
import sys
import tempfile
import importlib

import numpy as np
import soundfile as sf
import torch

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

PASS = 0
FAIL = 0


def check(label, condition, detail=""):
    global PASS, FAIL
    if condition:
        PASS += 1
        print(f"  [OK] {label}")
    else:
        FAIL += 1
        print(f"  [FAIL] {label} -- {detail}")


def test_global_matmul_precision():
    print("\n=== torch.set_float32_matmul_precision ===")
    import voice_clone  # noqa: F401 — importing sets precision
    # PyTorch doesn't expose a getter, but the call should not raise
    check("Module imported without error (precision set)", True)


def test_dtype_selection():
    print("\n=== Dtype selection ===")
    from voice_clone import _select_dtype

    dtype = _select_dtype()
    if torch.cuda.is_available():
        cap = torch.cuda.get_device_capability()
        if cap[0] >= 8:
            check("Ampere+ GPU -> bfloat16", dtype == torch.bfloat16,
                  f"got {dtype}")
        else:
            check("Pre-Ampere GPU -> float16", dtype == torch.float16,
                  f"got {dtype}")
    else:
        check("CPU -> float32", dtype == torch.float32, f"got {dtype}")


def test_attn_selection():
    print("\n=== Attention implementation selection ===")
    from voice_clone import _select_attn_impl

    impl = _select_attn_impl()
    try:
        import flash_attn  # noqa: F401
        check("flash-attn installed -> flash_attention_2",
              impl == "flash_attention_2", f"got {impl}")
    except ImportError:
        check("flash-attn not installed -> sdpa fallback",
              impl == "sdpa", f"got {impl}")


def test_generation_params():
    print("\n=== Generation parameters ===")
    from voice_clone import GENERATION_PARAMS

    check("top_k reduced to 20",
          GENERATION_PARAMS["top_k"] == 20,
          f"got {GENERATION_PARAMS['top_k']}")
    check("temperature reduced to 0.7",
          GENERATION_PARAMS["temperature"] == 0.7,
          f"got {GENERATION_PARAMS['temperature']}")
    check("max_new_tokens capped at 1024",
          GENERATION_PARAMS["max_new_tokens"] == 1024,
          f"got {GENERATION_PARAMS['max_new_tokens']}")
    check("subtalker_top_k reduced to 20",
          GENERATION_PARAMS["subtalker_top_k"] == 20,
          f"got {GENERATION_PARAMS['subtalker_top_k']}")
    check("subtalker_temperature reduced to 0.7",
          GENERATION_PARAMS["subtalker_temperature"] == 0.7,
          f"got {GENERATION_PARAMS['subtalker_temperature']}")


def test_ref_audio_trimming():
    print("\n=== Reference audio trimming ===")
    from voice_clone import _trim_reference_audio, REF_AUDIO_MAX_SEC

    # Create a 20-second test WAV
    sr = 24000
    long_audio = np.random.randn(int(20 * sr)).astype(np.float32) * 0.1
    with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as f:
        sf.write(f.name, long_audio, sr)
        long_path = f.name

    # Create a 4-second test WAV (under limit)
    short_audio = np.random.randn(int(4 * sr)).astype(np.float32) * 0.1
    with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as f:
        sf.write(f.name, short_audio, sr)
        short_path = f.name

    try:
        # Long audio should be trimmed
        trimmed = _trim_reference_audio(long_path)
        check("Long audio (20s) is trimmed", trimmed != long_path)
        if trimmed != long_path:
            info = sf.info(trimmed)
            check(f"Trimmed to <= {REF_AUDIO_MAX_SEC}s",
                  info.duration <= REF_AUDIO_MAX_SEC + 0.1,
                  f"got {info.duration:.1f}s")
            os.unlink(trimmed)

        # Short audio should NOT be trimmed
        result = _trim_reference_audio(short_path)
        check("Short audio (4s) is NOT trimmed", result == short_path)
    finally:
        os.unlink(long_path)
        os.unlink(short_path)


def test_model_mode_env():
    print("\n=== Model mode env var ===")
    from voice_clone import get_model_mode, BASE_MODEL_ID, CUSTOM_VOICE_MODEL_ID

    mode = get_model_mode()
    check("Default mode is 'base'", mode == "base", f"got '{mode}'")
    check("BASE_MODEL_ID is 1.7B", "1.7B" in BASE_MODEL_ID, BASE_MODEL_ID)
    check("CUSTOM_VOICE_MODEL_ID is 0.6B", "0.6B" in CUSTOM_VOICE_MODEL_ID,
          CUSTOM_VOICE_MODEL_ID)


def test_custom_voice_clone_blocked():
    print("\n=== CustomVoice clone attempt blocked ===")
    # Simulate QWEN_TTS_MODE=custom_voice by patching
    import voice_clone
    original = voice_clone._MODEL_MODE
    voice_clone._MODEL_MODE = "custom_voice"
    try:
        voice_clone.create_voice_profile("dummy.wav")
        check("Should have raised ValueError", False)
    except ValueError as e:
        check("create_voice_profile raises ValueError for custom_voice mode",
              "Base model" in str(e), str(e))
    except Exception as e:
        check("Unexpected error type", False, str(e))
    finally:
        voice_clone._MODEL_MODE = original


def test_tts_module_imports():
    print("\n=== TTS module structure ===")
    from tts import generate_audio_stream, split_into_chunks
    import inspect

    sig = inspect.signature(generate_audio_stream)
    params = list(sig.parameters.keys())
    check("generate_audio_stream has 'use_custom_voice' param",
          "use_custom_voice" in params, f"params: {params}")
    check("generate_audio_stream has 'custom_voice_speaker' param",
          "custom_voice_speaker" in params, f"params: {params}")
    check("generate_audio_stream has 'voice_profile_id' param",
          "voice_profile_id" in params, f"params: {params}")


def test_voice_profile_persistence():
    print("\n=== Voice profile persistence ===")
    from voice_clone import (
        VOICE_PROFILE_DIR, list_saved_profiles, load_default_profile,
        has_profile, save_profile_to_disk, load_profile_from_disk,
        _PROFILE_CACHE, _cache_lock,
    )
    import inspect

    check("VOICE_PROFILE_DIR exists", VOICE_PROFILE_DIR.exists(),
          str(VOICE_PROFILE_DIR))
    check("VOICE_PROFILE_DIR is named Voice_Profile",
          VOICE_PROFILE_DIR.name == "Voice_Profile")

    # Verify API signatures
    sig_create = inspect.signature(
        __import__('voice_clone').create_voice_profile
    )
    check("create_voice_profile accepts voice_name",
          "voice_name" in sig_create.parameters,
          f"params: {list(sig_create.parameters)}")

    check("list_saved_profiles returns a list",
          isinstance(list_saved_profiles(), list))

    # load_default_profile returns None when no profiles saved
    # (since Voice_Profile/ might have profiles from prior runs, just check type)
    result = load_default_profile()
    check("load_default_profile returns str or None",
          result is None or isinstance(result, str), f"got {type(result)}")

    # has_profile checks disk too
    check("has_profile('nonexistent') is False", has_profile("nonexistent") is False)


def test_app_default_profile():
    print("\n=== App default profile loading ===")
    source = open(
        os.path.join(os.path.dirname(os.path.dirname(__file__)), "app.py"),
        encoding="utf-8"
    ).read()

    check("App imports load_default_profile",
          "from voice_clone import load_default_profile" in source)
    check("App calls load_default_profile()",
          "_default_profile_id = load_default_profile()" in source)
    check("voice_profile_state initialized with default",
          "voice_profile_state = gr.State(_default_profile_id)" in source)
    check("create_voice_profile passes voice_name",
          "voice_name=v_name.strip()" in source)


if __name__ == "__main__":
    print("=" * 60)
    print("Latency Optimization Tests")
    print("=" * 60)

    test_global_matmul_precision()
    test_dtype_selection()
    test_attn_selection()
    test_generation_params()
    test_ref_audio_trimming()
    test_model_mode_env()
    test_custom_voice_clone_blocked()
    test_tts_module_imports()
    test_voice_profile_persistence()
    test_app_default_profile()

    print("\n" + "=" * 60)
    print(f"Results: {PASS} passed, {FAIL} failed")
    print("=" * 60)
    sys.exit(1 if FAIL > 0 else 0)