inference_example.py

#!/usr/bin/env python3
"""
Simple inference example for fine-tuned F5-TTS model
Shows how to use your trained checkpoint to generate speech
"""

import sys
sys.path.append('/home/husrcf/Code/AIAA/AIAA2205-assignment2-F5-TTS/')

import torch
import soundfile as sf
from src.f5_tts.infer.utils_infer import load_model, load_vocoder, infer_process

# ============================================================================
# Configuration
# ============================================================================

# Model settings
MODEL_NAME = "F5TTS_v1_Base"
CHECKPOINT_PATH = "./ckpts/sichuan_data/model_last.pt"  # Your trained checkpoint
VOCAB_PATH = "./data/sichuan_data/vocab.txt"  # Your custom vocabulary

# Reference audio (from your dataset or a new recording)
REF_AUDIO = "./data/dialect_data/speech/G0001_0009.wav"
REF_TEXT = "三国吴国顾雍是个什么样的人"  # Text spoken in reference audio

# Text to generate
GEN_TEXT = "你好，世界。这是一个测试。"  # Change this to what you want to synthesize

# Output path
OUTPUT_PATH = "./output_generated.wav"

# Generation settings
NFE_STEP = 32  # Number of function evaluations (higher = better quality, slower)
CFG_STRENGTH = 2.0  # Classifier-free guidance strength (1.0-3.0)
SWAY_SAMPLING_COEF = -1.0  # Sway sampling coefficient
SPEED = 1.0  # Speech speed multiplier

# Device
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"

# ============================================================================
# Main Inference
# ============================================================================

def main():
    print("=" * 70)
    print("F5-TTS Inference - Using Fine-tuned Model")
    print("=" * 70)

    print(f"\n📦 Model: {MODEL_NAME}")
    print(f"💾 Checkpoint: {CHECKPOINT_PATH}")
    print(f"📝 Vocab: {VOCAB_PATH}")
    print(f"🔊 Reference Audio: {REF_AUDIO}")
    print(f"💬 Reference Text: {REF_TEXT}")
    print(f"✍️  Text to Generate: {GEN_TEXT}")
    print(f"🎯 Device: {DEVICE}")
    print(f"⚙️  NFE Steps: {NFE_STEP}, CFG: {CFG_STRENGTH}, Speed: {SPEED}")

    print("\n" + "-" * 70)
    print("Loading model...")
    print("-" * 70)

    # Load the fine-tuned model
    try:
        model, _ = load_model(
            model_cls=MODEL_NAME,
            model_cfg=None,  # Will use default config
            ckpt_step=None,
            vocab_file=VOCAB_PATH,
            ckpt_file=CHECKPOINT_PATH,
            device=DEVICE
        )
        print("✓ Model loaded successfully!")
    except Exception as e:
        print(f"✗ Error loading model: {e}")
        return

    print("\n" + "-" * 70)
    print("Loading vocoder...")
    print("-" * 70)

    # Load vocoder (converts mel-spectrogram to audio)
    try:
        vocoder = load_vocoder(
            vocoder_name="vocos",
            is_local=False,
            local_path=None
        )
        print("✓ Vocoder loaded successfully!")
    except Exception as e:
        print(f"✗ Error loading vocoder: {e}")
        return

    print("\n" + "-" * 70)
    print("Generating speech...")
    print("-" * 70)

    # Generate speech
    try:
        generated_audio, sample_rate, spectrogram = infer_process(
            ref_audio=REF_AUDIO,
            ref_text=REF_TEXT,
            gen_text=GEN_TEXT,
            model_obj=model,
            vocoder=vocoder,
            mel_spec_type="vocos",
            show_info=print,  # Print progress
            progress=None,
            target_rms=0.1,
            cross_fade_duration=0.15,
            nfe_step=NFE_STEP,
            cfg_strength=CFG_STRENGTH,
            sway_sampling_coef=SWAY_SAMPLING_COEF,
            speed=SPEED,
            fix_duration=None,
            device=DEVICE
        )

        print("✓ Speech generated successfully!")

        # Save output
        sf.write(OUTPUT_PATH, generated_audio, sample_rate)

        duration = len(generated_audio) / sample_rate
        print(f"\n✓ Audio saved to: {OUTPUT_PATH}")
        print(f"  Duration: {duration:.2f}s")
        print(f"  Sample rate: {sample_rate}Hz")

        print("\n" + "=" * 70)
        print("Inference completed! 🎉")
        print("=" * 70)

    except Exception as e:
        print(f"\n✗ Error during generation: {e}")
        import traceback
        traceback.print_exc()


if __name__ == "__main__":
    main()