from omnivoice import OmniVoice
import torch
import torchaudio
import os

def main():
    print("Loading OmniVoice model (this might take a moment)...")
    model = OmniVoice.from_pretrained(
        "model_weights",
        device_map="cuda:0",
        dtype=torch.float16
    )
    
    # print("\n--- 1. AUTO VOICE (TỰ SINH GIỌNG) ---")
    # text_auto = "Chào bạn, đây là giọng nói tự sinh ra từ mô hình Omni Voice."
    # print(f"Sinh giọng tự động đọc câu: '{text_auto}'")
    # audio_auto = model.generate(
    #     text=text_auto,
    #     num_step=32
    # )
    # torchaudio.save("output_auto.wav", audio_auto[0], 24000)
    # print("Đã lưu -> output_auto.wav")


    print("\n--- 2. ZERO-SHOT VOICE CLONING (SAO CHÉP GIỌNG) ---")
    text_clone = """Đặc phái viên Mỹ kêu gọi lập liên minh mới thay thế NATO
Đặc phái viên Mỹ kêu gọi lập liên minh mới thay thế NATO
1
Tướng Kellogg, đặc phái viên của ông Trump về Ukraine, gợi ý lập liên minh mới thay thế NATO, có Kiev tham gia.

"NATO rốt cuộc tỏ ra hèn nhát. Có lẽ chúng ta cần một NATO mới, một cấu trúc phòng thủ mới", tướng về hưu Keith Kellogg, đặc phái viên của Tổng thống Mỹ Donald Trump về vấn đề Ukraine, nói trong phỏng vấn ngày 2/4 với FOX News.

Tướng Kellogg chỉ trích các đồng minh trong NATO vì đã không triển khai lực lượng hỗ trợ Mỹ trong chiến sự Iran, lặp lại quan điểm của Tổng thống Trump. Ông Trump trước đó cũng cho rằng NATO "hèn nhát, hổ giấy" vì không giúp Mỹ mở lại eo biển Hormuz, đồng thời dọa rút Mỹ khỏi liên minh quân sự lớn nhất thế giới."""
    ref_path = "ref.wav" # <-- Tên file âm thanh mẫu (khoảng 3-10 giây giọng người cần clone)

    if os.path.exists(ref_path):
        print(f"Tìm thấy '{ref_path}'. Đang tiến hành clone giọng...")
        ref_text_chuot = "trường quê em do bố của em xây kỹ nên sạch sẽ đẹp đẽ và rất kiên cố mọi người ơi" # <-- Sửa theo nội dung file ref.wav của bạn

        audio_clone = model.generate(
            text=text_clone,
            ref_audio=ref_path,
            ref_text=ref_text_chuot,  # Truyền chữ mẫu vào để model khớp giọng tốt nhất
            num_step=32,
            # speed=1.2,
            # guidance_scale=2.0,
            # t_shift=0.1,
            # layer_penalty_factor=5.0,
            # position_temperature=5.0,
            # class_temperature=0.0,
            # denoise=True,
            # preprocess_prompt=True,
            # postprocess_output=True,
            # audio_chunk_duration=15.0,
            # audio_chunk_threshold=30.0,
        )
        torchaudio.save("output_clone.wav", audio_clone[0], 24000)
        print("Đã lưu giọng clone -> output_clone.wav")
    else:
        print(f"⚠️ KHÔNG TÌM THẤY file âm thanh mẫu '{ref_path}'!")
        print(f"Để thử tính năng Voice Clone, bạn hãy:")
        print(f"  1. Chuẩn bị 1 file âm thanh (định dạng wav/mp3) chứa giọng nói muốn copy (dài khoảng 3 đến 10 giây).")
        print(f"  2. Để nó vào thư mục này (/home/thien/chum/OmniVoice/) và đổi tên thành 'ref.wav'")
        print(f"  3. Chạy lại lệnh 'python test_tts.py' để xem kết quả!")

if __name__ == "__main__":
    main()