import torch
import soundfile as sf
import gradio as gr
from transformers import AutoModel, AutoProcessor

# ==========================================
# Audio8 Model
# ==========================================

MODEL_ID = "Audio8/Audio8-TTS-Preview-0.6b"

# Select GPU if available
device = "cuda" if torch.cuda.is_available() else "cpu"

# Use BF16 on GPU, FP32 on CPU
dtype = torch.bfloat16 if device == "cuda" else torch.float32

print("=" * 60)
print("Audio8 Text-to-Speech")
print("=" * 60)

print("PyTorch:", torch.__version__)
print("CUDA Available:", torch.cuda.is_available())

if torch.cuda.is_available():
    print("GPU:", torch.cuda.get_device_name(0))

# ==========================================
# Load Processor
# ==========================================

print("\nLoading processor...")

processor = AutoProcessor.from_pretrained(
    MODEL_ID,
    trust_remote_code=True
)

# ==========================================
# Load Model
# ==========================================

print("Loading Audio8 model...")

model = AutoModel.from_pretrained(
    MODEL_ID,
    trust_remote_code=True,
    dtype=dtype
).eval().to(device)

print("✅ Audio8 model loaded successfully!")

# ==========================================
# Generate Speech
# ==========================================

def generate_speech(text):

    # Check empty input
    if not text or not text.strip():
        return None

    print("\n" + "=" * 50)
    print("Generating speech...")
    print("Text:", text)

    # Process text
    inputs = processor(
        text=[text],
        return_tensors="pt"
    )

    # Move inputs to GPU/CPU
    inputs = {
        name: value.to(device)
        for name, value in inputs.items()
    }

    # Generate speech tokens
    with torch.inference_mode():

        output = model.generate(
            **inputs,
            max_new_tokens=1024,
            temperature=0.8,
            top_p=0.95,
            top_k=50,
            do_sample=True,
            return_dict_in_generate=True,
        )

    # Decode generated audio
    waveforms, waveform_lengths = model.decode_audio(
        output.codes
    )

    # Convert to NumPy
    audio = waveforms[
        0,
        :int(waveform_lengths[0])
    ].float().cpu().numpy()

    # ==========================================
    # Normalize Audio Volume
    # ==========================================

    peak = abs(audio).max()

    if peak > 0:
        audio = audio / peak
        audio = audio * 0.95

    # ==========================================
    # Save Audio
    # ==========================================

    output_file = "output.wav"

    sf.write(
        output_file,
        audio,
        model.config.codec_sample_rate
    )

    print("✅ Speech generated!")
    print("🎵 Saved:", output_file)

    return output_file


# ==========================================
# Gradio User Interface
# ==========================================

with gr.Blocks(title="SmartVoiceAI") as app:

    gr.Markdown(
        """
        # 🔊 SmartVoiceAI

        ### AI Text-to-Speech

        Enter your text below and generate natural-sounding speech
        using the Audio8 TTS model.
        """
    )

    # Text input
    text_input = gr.Textbox(
        label="Enter Text",
        placeholder="Type something here...",
        lines=6
    )

    # Generate button
    generate_button = gr.Button(
        "🔊 Generate Speech",
        variant="primary"
    )

    # Audio output
    audio_output = gr.Audio(
        label="Generated Voice",
        type="filepath",
        autoplay=True
    )

    # Button action
    generate_button.click(
        fn=generate_speech,
        inputs=text_input,
        outputs=audio_output
    )

    gr.Markdown(
        """
        ---
        **Model:** Audio8-TTS-Preview-0.6b  
        **Device:** CUDA / CPU
        """
    )


# ==========================================
# Start Application
# ==========================================

if __name__ == "__main__":

    app.launch()