import spaces
import gradio as gr
import json
import torch
import wavio
from tqdm import tqdm
from huggingface_hub import snapshot_download
from models import AudioDiffusion, DDPMScheduler
from audioldm.audio.stft import TacotronSTFT
from audioldm.variational_autoencoder import AutoencoderKL
from pydub import AudioSegment
from gradio import Markdown

import torch
#from diffusers.models.autoencoder_kl import AutoencoderKL
from diffusers import DiffusionPipeline,AudioPipelineOutput
from transformers import CLIPTextModel, T5EncoderModel, AutoModel, T5Tokenizer, T5TokenizerFast
from typing import Union
from diffusers.utils.torch_utils import randn_tensor
from tqdm import tqdm
from TangoFlux import TangoFluxInference



tangoflux = TangoFluxInference(path="declare-lab/TangoFlux")



@spaces.GPU(duration=15)
def gradio_generate(prompt, output_format, steps, guidance,duration=10):

    output_wave = tangoflux.generate(prompt,steps=steps,guidance=guidance,duration=duration)
    output_wave = pipe(prompt,steps,guidance) ## Using pipeliine automatically uses flash attention for torch2.0 above
    #output_wave = tango.generate(prompt, steps, guidance)
    # output_filename = f"{prompt.replace(' ', '_')}_{steps}_{guidance}"[:250] + ".wav"
    output_wave = output_wave.audios[0]
    output_filename = "temp.wav"
    wavio.write(output_filename, output_wave, rate=16000, sampwidth=2)

    if (output_format == "mp3"):
        AudioSegment.from_wav("temp.wav").export("temp.mp3", format = "mp3")
        output_filename = "temp.mp3"

    return output_filename

description_text = """
<p><a href="https://huggingface.co/spaces/declare-lab/tango2/blob/main/app.py?duplicate=true"> <img style="margin-top: 0em; margin-bottom: 0em" src="https://bit.ly/3gLdBN6" alt="Duplicate Space"></a> For faster inference without waiting in queue, you may duplicate the space and upgrade to a GPU in the settings. <br/><br/>
Generate audio using Tango2 by providing a text prompt. Tango2 was built from Tango and was trained on <a href="https://huggingface.co/datasets/declare-lab/audio-alpaca">Audio-alpaca</a>
<br/><br/> This is the demo for Tango2 for text to audio generation: <a href="https://arxiv.org/abs/2404.09956">Read our paper.</a>
<p/>
"""
# Gradio input and output components
input_text = gr.Textbox(lines=2, label="Prompt")
output_format = gr.Radio(label = "Output format", info = "The file you can dowload", choices = ["mp3", "wav"], value = "wav")
output_audio = gr.Audio(label="Generated Audio", type="filepath")
denoising_steps = gr.Slider(minimum=10, maximum=100, value=25, step=1, label="Steps", interactive=True)
guidance_scale = gr.Slider(minimum=1, maximum=10, value=3, step=0.1, label="Guidance Scale", interactive=True)
duration_scale = gr.Slider(minimum=1, maximum=30, value=10, step=1, label="Duration", interactive=True)

# Gradio interface
gr_interface = gr.Interface(
    fn=gradio_generate,
    inputs=[input_text, output_format, denoising_steps, guidance_scale,duration_scale],
    outputs=[output_audio],
    title="TangoFlux: Aligning Diffusion-based Text-to-Audio Generations through Direct Preference Optimization",
    description=description_text,
    allow_flagging=False,
    examples=[
        ["Quiet speech and then and airplane flying away"],
        ["A bicycle peddling on dirt and gravel followed by a man speaking then laughing"],
        ["Ducks quack and water splashes with some animal screeching in the background"],
        ["Describe the sound of the ocean"],
        ["A woman and a baby are having a conversation"],
        ["A man speaks followed by a popping noise and laughter"],
        ["A cup is filled from a faucet"],
        ["An audience cheering and clapping"],
        ["Rolling thunder with lightning strikes"],
        ["A dog barking and a cat mewing and a racing car passes by"],
        ["Gentle water stream, birds chirping and sudden gun shot"],
        ["A man talking followed by a goat baaing then a metal gate sliding shut as ducks quack and wind blows into a microphone."],
        ["A dog barking"],
        ["A cat meowing"],
        ["Wooden table tapping sound while water pouring"],
        ["Applause from a crowd with distant clicking and a man speaking over a loudspeaker"],
        ["two gunshots followed by birds flying away while chirping"],
        ["Whistling with birds chirping"],
        ["A person snoring"],
        ["Motor vehicles are driving with loud engines and a person whistles"],
        ["People cheering in a stadium while thunder and lightning strikes"],
        ["A helicopter is in flight"],
        ["A dog barking and a man talking and a racing car passes by"],
    ],
    cache_examples="lazy", # Turn on to cache.
)

# Launch Gradio app
gr_interface.queue(10).launch()