import time
from omnivoice import OmniVoice
import soundfile as sf
import torch

model = OmniVoice.from_pretrained(
    "/home/riftuser/OmniVoice/exp_v1/omnivoice_finetune/checkpoint-500",
    device_map="cuda:0",
    dtype=torch.float16
)

saudi_tts_text = """
السلام عليكم يا شباب، كيف الحال؟ 

اليوم عندي تقريبًا 3 meetings مهمة، وأول meeting بتبدأ الساعة 10:30 الصباح. [sigh] [sigh]

بصراحة كنت ناوي أخلص الـ report بدري، لكن الـ internet صار بطيء بشكل مو طبيعي. [dissatisfaction-hnn] [sigh] [dissatisfaction-hnn]

قلت خلاص، خلني آخذ coffee وأروق شوي قبل ما أبدأ الشغل. [laughter] [laughter] [confirmation-en]

وبعدين اكتشفت إن الـ laptop يحتاج update من أمس! [surprise-oh] [dissatisfaction-hnn]

قلت يا ساتر، شكله يوم طويل جدًا. [sigh] [laughter]

لكن الحمد لله الأمور مشت تمام بالنهاية. 
"""

# Measure generation time
start_time = time.time()

audio = model.generate(
    text=saudi_tts_text,
    ref_audio="ref_audio/women_ref_1.mp3",
    ref_text="شوفي يا حلوة هالكريم الجديد للبشرة، يخلي وجهك مثل القمر! ",
    instruct = "female, young adult, high pitch",
    speed = 1.1,
    num_step = 25,
    guidance_scale=2.0,
    t_shift=0.1,
    position_temperature=3,
    layer_penalty_factor=5.0,
    
)

generation_time = time.time() - start_time

# Save audio
sf.write("out_1.wav", audio[0], 24000)

# Calculate audio duration
audio_duration = len(audio[0]) / 24000

# Calculate RTF
rtf = generation_time / audio_duration

print(f"Generation Time: {generation_time:.2f} sec")
print(f"Audio Duration: {audio_duration:.2f} sec")
print(f"RTF: {rtf:.4f}")