import time
import torch
import soundfile as sf
from qwen_tts import Qwen3TTSModel

# Load the model
model = Qwen3TTSModel.from_pretrained(
    "Qwen/Qwen3-TTS-12Hz-0.6B-CustomVoice",
    device_map="cuda:0",
    dtype=torch.bfloat16,
)
print(model.get_supported_languages())

with open("input.txt", "r") as f:
    text = f.read()
    f.close()

# Generate speech with specific instructions
start_time = time.perf_counter()
wavs, sr = model.generate_custom_voice(
    text=text,
    language="English", 
    speaker="Aiden",
    instruct="read the text in a calm and soothing voice, with a slight emphasis on key points, and maintain a steady pace throughout the narration.", 
)
end_time = time.perf_counter()

# Save the generated audio
sf.write("output_custom_voice.wav", wavs[0], sr)

# Print timing
duration = end_time - start_time
audio_duration = len(wavs[0]) / sr
print(f"Generation time: {duration:.2f}s")
print(f"Audio duration: {audio_duration:.2f}s")
print(f"Real-time factor: {duration / audio_duration:.2f}x")