from __future__ import annotations

import gradio as gr
try:
    import spaces
except ImportError:  # 'spaces' is only provided by Hugging Face Spaces
    import types as _types

    def _gpu(*args, **kwargs):
        if len(args) == 1 and callable(args[0]) and not kwargs:
            return args[0]

        def _decorator(func):
            return func

        return _decorator

    spaces = _types.SimpleNamespace(GPU=_gpu)

from pyharp import *


try:  # torch>=2.6 flipped torch.load(weights_only) to True; legacy ckpts need False
    import torch as _torch

    if getattr(_torch.load, "__harp_compat__", False) is False:
        _torch_load_orig = _torch.load

        def _torch_load_compat(*args, **kwargs):
            kwargs.setdefault("weights_only", False)
            return _torch_load_orig(*args, **kwargs)

        _torch_load_compat.__harp_compat__ = True
        _torch.load = _torch_load_compat
except Exception:  # torch not installed / unexpected API -- nothing to patch
    pass


from bark import generate_audio, preload_models, SAMPLE_RATE
from scipy.io.wavfile import write as write_wav
import tempfile
import os

# Download and load all models once
preload_models()


model_card = ModelCard(
    name="Bark",
    description="Bark is a transformer-based text-to-audio model created by Suno. Bark can generate highly realistic, multilingual speech as well as other audio - including music, background noise and simple sound effects. The model can also produce nonverbal communications like laughing, sighing and crying.",
    author="suno-ai",
    tags=["text-to-audio", "speech-synthesis", "multilingual", "music-generation"],
)


@spaces.GPU
def process_fn(text_prompt, history_prompt, text_temp, waveform_temp):
    history_prompt_val = history_prompt if history_prompt != "None" else None

    # Generate audio from text
    audio_array = generate_audio(
        text_prompt,
        history_prompt=history_prompt_val,
        text_temp=text_temp,
        waveform_temp=waveform_temp,
        silent=True, # Disable progress bar for Gradio app
        output_full=False # Only return the audio array
    )

    # Save audio to a temporary WAV file
    output_file = tempfile.NamedTemporaryFile(suffix=".wav", delete=False)
    output_file.close()
    write_wav(output_file.name, SAMPLE_RATE, audio_array)

    return output_file.name


with gr.Blocks() as demo:
    input_components = [
        gr.Textbox(label="Text Prompt", info="The text to be converted into audio. Bark supports multilingual speech, music, and nonverbal communications.").harp_required(True),
        gr.Dropdown(choices=["None", "announcer", "en_speaker_0", "en_speaker_1", "en_speaker_2", "en_speaker_3", "en_speaker_4", "en_speaker_5", "en_speaker_6", "en_speaker_7", "en_speaker_8", "en_speaker_9"], value="en_speaker_0", label="Voice Prompt (Speaker)", info="Choose a speaker to use as a voice prompt for audio cloning. 'None' will use a default voice."),
        gr.Slider(minimum=0.0, maximum=1.0, step=0.05, value=0.7, label="Text Generation Temperature", info="Controls the diversity of text token generation (1.0 more diverse, 0.0 more conservative)."),
        gr.Slider(minimum=0.0, maximum=1.0, step=0.05, value=0.7, label="Waveform Generation Temperature", info="Controls the diversity of waveform token generation (1.0 more diverse, 0.0 more conservative)."),
    ]
    output_components = [
        gr.Audio(type="filepath", label="Generated Audio"),
    ]
    build_endpoint(
        model_card=model_card,
        input_components=input_components,
        output_components=output_components,
        process_fn=process_fn,
    )

demo.queue().launch(share=True, show_error=False, pwa=True)